Upload DeFTAN2.py

DeFTAN2.py

@@ -0,0 +1,329 @@
+import math
+from collections import OrderedDict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from packaging.version import parse as V
+from torch.nn import init
+from torch.nn.parameter import Parameter
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+
+class DeFTAN2(nn.Module):
+    def __init__(self, n_srcs=1, win=512, n_mics=4, n_layers=6, att_dim=64, hidden_dim=256, n_head=4, emb_dim=64, emb_ks=4, emb_hs=1, dropout=0.1, eps=1.0e-5):
+        super().__init__()
+        self.n_srcs = n_srcs
+        self.win = win
+        self.hop = win // 2
+        self.n_layers = n_layers
+        self.n_mics = n_mics
+        self.emb_dim = emb_dim
+        assert win % 2 == 0
+
+        t_ksize = 3
+        ks, padding = (t_ksize, 3), (t_ksize // 2, 1)
+        self.up_conv = nn.Sequential(
+            nn.Conv2d(2 * n_mics, emb_dim * n_head, ks, padding=padding),
+            nn.GroupNorm(1, emb_dim * n_head, eps=eps),
+            SDB2d(emb_dim * n_head, emb_dim, n_head)
+        )
+        self.blocks = nn.ModuleList([])
+        for idx in range(n_layers):
+            self.blocks.append(DeFTAN2block(idx, emb_dim, emb_ks, emb_hs, att_dim, hidden_dim, n_head, dropout, eps))
+        self.down_conv = nn.Sequential(
+            nn.Conv2d(emb_dim, 2 * n_srcs * n_head, ks, padding=padding),
+            SDB2d(2 * n_srcs * n_head, 2 * n_srcs, n_head))
+
+    def pad_signal(self, input):
+        # input is the waveforms: (B, T) or (B, 1, T)
+        # reshape and padding
+        if input.dim() not in [2, 3]:
+            raise RuntimeError("Input can only be 2 or 3 dimensional.")
+
+        if input.dim() == 2:
+            input = input.unsqueeze(1)
+        batch_size = input.size(0)
+        nchannel = input.size(1)
+        nsample = input.size(2)
+        rest = self.win - (self.hop + nsample % self.win) % self.win
+        if rest > 0:
+            pad = Variable(torch.zeros(batch_size, nchannel, rest)).type(input.type())
+            input = torch.cat([input, pad], 2)
+
+        pad_aux = Variable(torch.zeros(batch_size, nchannel, self.hop)).type(input.type())
+        input = torch.cat([pad_aux, input, pad_aux], 2)
+
+        return input, rest
+
+    def forward(self, input: Union[torch.Tensor]) -> Tuple[List[Union[torch.Tensor]], torch.Tensor, OrderedDict]:
+        input, rest = self.pad_signal(input)
+        B, M, N = input.size()                                                      # batch B, mic M, time samples N
+        mix_std_ = torch.std(input, dim=(1, 2), keepdim=True)  # [B, 1, 1]
+        input = input / mix_std_  # RMS normalization
+        # Encoding
+        stft_input = torch.stft(input.view([-1, N]), n_fft=self.win, hop_length=self.hop, window=torch.hann_window(self.win).type(input.type()), return_complex=False)
+        _, F, T, _ = stft_input.size()                                              # B*M , F= num freqs, T= num frame, 2= real imag
+        xi = stft_input.view([B, M, F, T, 2])                                       # B*M, F, T, 2 -> B, M, F, T, 2
+        xi = xi.permute(0, 1, 4, 3, 2).contiguous()                                 # [B, M, 2, T, F]
+        xi = xi.view([B, M * 2, T, F])                                              # [B, 2*M, T, F]
+        # Separation
+        feature = self.up_conv(xi)                                                    # [B, C, T, F]
+        for ii in range(self.n_layers):
+            feature = self.blocks[ii](feature)                                          # [B, C, T, F]
+        xo = self.down_conv(feature).view([B, self.n_srcs, 2, T, F]).view([B * self.n_srcs, 2, T, F])
+        # Decoding
+        xo = xo.permute(0, 3, 2, 1).type(input.type())                        # [B*n_srcs, 2, T, F] -> [B*n_srcs, F, T, 2]
+        istft_input = torch.complex(xo[:, :, :, 0], xo[:, :, :, 1])
+        istft_output = torch.istft(istft_input, n_fft=self.win, hop_length=self.hop, window=torch.hann_window(self.win).type(input.type()), return_complex=False)
+
+        output = istft_output[:, self.hop:-(rest + self.hop)].unsqueeze(1)          # [B*n_srcs, 1, N]
+        output = output.view([B, self.n_srcs, -1])                                  # [B, n_srcs, N]
+        output = output * mix_std_  # reverse the RMS normalization
+
+        return output
+
+
+class SDB1d(nn.Module):
+    def __init__(self, in_channels, out_channels, groups):
+        super().__init__()
+        assert in_channels // out_channels == groups
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.groups = groups
+        self.blocks = nn.ModuleList([])
+        for idx in range(groups):
+            self.blocks.append(nn.Sequential(
+                nn.Conv1d(out_channels * ((idx > 0) + 1), out_channels, kernel_size=3, padding=1),
+                nn.GroupNorm(1, out_channels, 1e-5),
+                nn.PReLU(out_channels)
+            ))
+
+    def forward(self, x):
+        B, C, L = x.size()
+        g = self.groups
+        # x = x.view(B, g, C//g, L).transpose(1, 2).reshape(B, C, L)
+        skip = x[:, ::g, :]
+        for idx in range(g):
+            output = self.blocks[idx](skip)
+            skip = torch.cat([output, x[:, idx+1::g, :]], dim=1)
+        return output
+
+
+class SDB2d(nn.Module):
+    def __init__(self, in_channels, out_channels, groups):
+        super().__init__()
+        assert in_channels // out_channels == groups
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.groups = groups
+        self.blocks = nn.ModuleList([])
+        for idx in range(groups):
+            self.blocks.append(nn.Sequential(
+                nn.Conv2d(out_channels * ((idx > 0) + 1), out_channels, kernel_size=(3, 3), padding=(1, 1)),
+                nn.GroupNorm(1, out_channels, 1e-5),
+                nn.PReLU(out_channels)
+            ))
+
+    def forward(self, x):
+        B, C, T, Q = x.size()
+        g = self.groups
+        # x = x.view(B, g, C//g, T, Q).transpose(1, 2).reshape(B, C, T, Q)
+        skip = x[:, ::g, :, :]
+        for idx in range(g):
+            output = self.blocks[idx](skip)
+            skip = torch.cat([output, x[:, idx+1::g, :, :]], dim=1)
+        return output
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+
+class CEA(nn.Module):
+    def __init__(self, dim, heads, dim_head, dropout):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.cv_qk = nn.Sequential(
+            nn.Conv1d(dim, dim * 2, kernel_size=3, padding=1, bias = False),
+            nn.GLU(dim=1))
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_k = nn.Linear(dim, inner_dim, bias=False)
+        self.to_v = nn.Linear(dim, inner_dim, bias = False)
+
+        self.att_drop = nn.Dropout(dropout)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        qk = self.cv_qk(x.transpose(1, 2)).transpose(1, 2)
+        q = rearrange(self.to_q(qk), 'b n (h d) -> b h n d', h = self.heads)
+        k = rearrange(self.to_k(qk), 'b n (h d) -> b h n d', h=self.heads)
+        v = rearrange(self.to_v(x), 'b n (h d) -> b h n d', h = self.heads)
+
+        weight = torch.matmul(F.softmax(k, dim=2).transpose(-1, -2), v) * self.scale
+        out = torch.matmul(F.softmax(q, dim=3), self.att_drop(weight))
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class DPFN(nn.Module):
+    def __init__(self, dim, hidden_dim, idx, dropout):
+        super().__init__()
+        self.PW1 = nn.Sequential(
+            nn.Linear(dim, hidden_dim//2),
+            nn.GELU(),
+	        nn.Dropout(dropout)
+        )
+        self.PW2 = nn.Sequential(
+            nn.Linear(dim, hidden_dim//2),
+            nn.GELU(),
+	        nn.Dropout(dropout)
+        )
+        self.DW_Conv = nn.Sequential(
+            nn.Conv1d(hidden_dim//2, hidden_dim//2, kernel_size=5, dilation=2**idx, padding='same'),
+            nn.GroupNorm(1, hidden_dim//2, 1e-5),
+            nn.PReLU(hidden_dim//2)
+        )
+        self.PW3 = nn.Sequential(
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        ffw_out = self.PW1(x)
+        dw_out = self.DW_Conv(self.PW2(x).transpose(1, 2)).transpose(1, 2)
+        out = self.PW3(torch.cat((ffw_out, dw_out), dim=2))
+        return out
+
+
+class DeFTAN2block(nn.Module):
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __init__(self, idx, emb_dim, emb_ks, emb_hs, att_dim, hidden_dim, n_head, dropout, eps):
+        super().__init__()
+        in_channels = emb_dim * emb_ks
+        self.F_norm = LayerNormalization4D(emb_dim, eps)
+        self.F_inv = SDB1d(in_channels, emb_dim, emb_ks)
+        self.F_att = PreNorm(emb_dim, CEA(emb_dim, n_head, att_dim, dropout))
+        self.F_ffw = PreNorm(emb_dim, DPFN(emb_dim, hidden_dim, idx, dropout))
+        self.F_linear = nn.ConvTranspose1d(emb_dim, emb_dim, emb_ks, stride=emb_hs)
+
+        self.T_norm = LayerNormalization4D(emb_dim, eps)
+        self.T_inv = SDB1d(in_channels, emb_dim, emb_ks)
+        self.T_att = PreNorm(emb_dim, CEA(emb_dim, n_head, att_dim, dropout))
+        self.T_ffw = PreNorm(emb_dim, DPFN(emb_dim, hidden_dim, idx, dropout))
+        self.T_linear = nn.ConvTranspose1d(emb_dim, emb_dim, emb_ks, stride=emb_hs)
+
+        self.emb_dim = emb_dim
+        self.emb_ks = emb_ks
+        self.emb_hs = emb_hs
+        self.n_head = n_head
+
+    def forward(self, x):
+        B, C, old_T, old_Q = x.shape
+        T = math.ceil((old_T - self.emb_ks) / self.emb_hs) * self.emb_hs + self.emb_ks
+        Q = math.ceil((old_Q - self.emb_ks) / self.emb_hs) * self.emb_hs + self.emb_ks
+        x = F.pad(x, (0, Q - old_Q, 0, T - old_T))
+
+        # F-transformer
+        input_ = x
+        F_feat = self.F_norm(input_)  # [B, C, T, Q]
+        F_feat = F_feat.transpose(1, 2).contiguous().view(B * T, C, Q)  # [BT, C, Q]
+        F_feat = F.unfold(F_feat[..., None], (self.emb_ks, 1), stride=(self.emb_hs, 1))  # [BT, C*emb_ks, -1]
+        F_feat = self.F_inv(F_feat)   # [BT, C, -1]
+
+        F_feat = F_feat.transpose(1, 2)  # [BT, -1, C]
+        F_feat = self.F_att(F_feat) + F_feat
+        F_feat = self.F_ffw(F_feat) + F_feat
+        F_feat = F_feat.transpose(1, 2)  # [BT, H, -1]
+
+        F_feat = self.F_linear(F_feat)  # [BT, C, Q]
+        F_feat = F_feat.view([B, T, C, Q])
+        F_feat = F_feat.transpose(1, 2).contiguous()  # [B, C, T, Q]
+        F_feat = F_feat + input_  # [B, C, T, Q]
+
+        # T-transformer
+        input_ = F_feat
+        T_feat = self.T_norm(input_)  # [B, C, T, F]
+        T_feat = T_feat.permute(0, 3, 1, 2).contiguous().view(B * Q, C, T)  # [BF, C, T]
+        T_feat = F.unfold(T_feat[..., None], (self.emb_ks, 1), stride=(self.emb_hs, 1))  # [BF, C*emb_ks, -1]
+        T_feat = self.T_inv(T_feat)   # [BF, C, -1]
+
+        T_feat = T_feat.transpose(1, 2)  # [BF, -1, C]
+        T_feat = self.T_att(T_feat) + T_feat
+        T_feat = self.T_ffw(T_feat) + T_feat
+        T_feat = T_feat.transpose(1, 2)  # [BF, H, -1]
+
+        T_feat = self.T_linear(T_feat)  # [BF, C, T]
+        T_feat = T_feat.view([B, Q, C, T])
+        T_feat = T_feat.permute(0, 2, 3, 1).contiguous()  # [B, C, T, Q]
+        T_feat = T_feat + input_  # [B, C, T, Q]
+
+        return T_feat
+
+
+class LayerNormalization4D(nn.Module):
+    def __init__(self, input_dimension, eps=1e-5):
+        super().__init__()
+        param_size = [1, input_dimension, 1, 1]
+        self.gamma = Parameter(torch.Tensor(*param_size).to(torch.float32))
+        self.beta = Parameter(torch.Tensor(*param_size).to(torch.float32))
+        init.ones_(self.gamma)
+        init.zeros_(self.beta)
+        self.eps = eps
+
+    def forward(self, x):
+        if x.ndim == 4:
+            _, C, _, _ = x.shape
+            stat_dim = (1,)
+        else:
+            raise ValueError("Expect x to have 4 dimensions, but got {}".format(x.ndim))
+        mu_ = x.mean(dim=stat_dim, keepdim=True)  # [B,1,T,F]
+        std_ = torch.sqrt(
+            x.var(dim=stat_dim, unbiased=False, keepdim=True) + self.eps
+        )  # [B,1,T,F]
+        x_hat = ((x - mu_) / std_) * self.gamma + self.beta
+        return x_hat
+
+
+class LayerNormalization4DCF(nn.Module):
+    def __init__(self, input_dimension, eps=1e-5):
+        super().__init__()
+        assert len(input_dimension) == 2
+        param_size = [1, input_dimension[0], 1, input_dimension[1]]
+        self.gamma = Parameter(torch.Tensor(*param_size).to(torch.float32))
+        self.beta = Parameter(torch.Tensor(*param_size).to(torch.float32))
+        init.ones_(self.gamma)
+        init.zeros_(self.beta)
+        self.eps = eps
+
+    def forward(self, x):
+        if x.ndim == 4:
+            stat_dim = (1, 3)
+        else:
+            raise ValueError("Expect x to have 4 dimensions, but got {}".format(x.ndim))
+        mu_ = x.mean(dim=stat_dim, keepdim=True)  # [B,1,T,1]
+        std_ = torch.sqrt(
+            x.var(dim=stat_dim, unbiased=False, keepdim=True) + self.eps
+        )  # [B,1,T,F]
+        x_hat = ((x - mu_) / std_) * self.gamma + self.beta
+        return x_hat