Upload model.py

model.py

@@ -0,0 +1,205 @@
+# coding: utf-8
+# Author：WangTianRui
+# Date ：2020/11/3 16:49
+
+import torch.nn as nn
+import torch
+from utils.conv_stft import *
+from utils.complexnn import *
+
+
+class DCCRN(nn.Module):
+    def __init__(self,
+                 rnn_layer=2, rnn_hidden=256,
+                 win_len=400, hop_len=100, fft_len=512, win_type='hann',
+                 use_clstm=True, use_cbn=False, masking_mode='E',
+                 kernel_size=5, kernel_num=(32, 64, 128, 256, 256, 256)
+                 ):
+        super(DCCRN, self).__init__()
+        self.rnn_layer = rnn_layer
+        self.rnn_hidden = rnn_hidden
+
+        self.win_len = win_len
+        self.hop_len = hop_len
+        self.fft_len = fft_len
+        self.win_type = win_type
+
+        self.use_clstm = use_clstm
+        self.use_cbn = use_cbn
+        self.masking_mode = masking_mode
+
+        self.kernel_size = kernel_size
+        self.kernel_num = (2,) + kernel_num
+
+        self.stft = ConvSTFT(self.win_len, self.hop_len, self.fft_len, self.win_type, 'complex', fix=True)
+        self.istft = ConviSTFT(self.win_len, self.hop_len, self.fft_len, self.win_type, 'complex', fix=True)
+
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        for idx in range(len(self.kernel_num) - 1):
+            self.encoder.append(
+                nn.Sequential(
+                    ComplexConv2d(
+                        self.kernel_num[idx],
+                        self.kernel_num[idx + 1],
+                        kernel_size=(self.kernel_size, 2),
+                        stride=(2, 1),
+                        padding=(2, 1)
+                    ),
+                    nn.BatchNorm2d(self.kernel_num[idx + 1]) if not use_cbn else ComplexBatchNorm(
+                        self.kernel_num[idx + 1]),
+                    nn.PReLU()
+                )
+            )
+        hidden_dim = self.fft_len // (2 ** (len(self.kernel_num)))
+
+        if self.use_clstm:
+            rnns = []
+            for idx in range(rnn_layer):
+                rnns.append(
+                    NavieComplexLSTM(
+                        input_size=hidden_dim * self.kernel_num[-1] if idx == 0 else self.rnn_hidden,
+                        hidden_size=self.rnn_hidden,
+                        batch_first=False,
+                        projection_dim=hidden_dim * self.kernel_num[-1] if idx == rnn_layer - 1 else None
+                    )
+                )
+                self.enhance = nn.Sequential(*rnns)
+        else:
+            self.enhance = nn.LSTM(
+                input_size=hidden_dim * self.kernel_num[-1],
+                hidden_size=self.rnn_hidden,
+                num_layers=2,
+                dropout=0.0,
+                batch_first=False
+            )
+            self.transform = nn.Linear(self.rnn_hidden, hidden_dim * self.kernel_num[-1])
+        for idx in range(len(self.kernel_num) - 1, 0, -1):
+            if idx != 1:
+                self.decoder.append(
+                    nn.Sequential(
+                        ComplexConvTranspose2d(
+                            self.kernel_num[idx] * 2,
+                            self.kernel_num[idx - 1],
+                            kernel_size=(self.kernel_size, 2),
+                            stride=(2, 1),
+                            padding=(2, 0),
+                            output_padding=(1, 0)
+                        ),
+                        nn.BatchNorm2d(self.kernel_num[idx - 1]) if not use_cbn else ComplexBatchNorm(
+                            self.kernel_num[idx - 1]),
+                        nn.PReLU()
+                    )
+                )
+            else:
+                self.decoder.append(
+                    nn.Sequential(
+                        ComplexConvTranspose2d(
+                            self.kernel_num[idx] * 2,
+                            self.kernel_num[idx - 1],
+                            kernel_size=(self.kernel_size, 2),
+                            stride=(2, 1),
+                            padding=(2, 0),
+                            output_padding=(1, 0)
+                        )
+                    )
+                )
+        if isinstance(self.enhance, nn.LSTM):
+            self.enhance.flatten_parameters()
+
+    def forward(self, x):
+        stft = self.stft(x)
+        # print("stft:", stft.size())
+        real = stft[:, :self.fft_len // 2 + 1]
+        imag = stft[:, self.fft_len // 2 + 1:]
+        # print("real imag:", real.size(), imag.size())
+        spec_mags = torch.sqrt(real ** 2 + imag ** 2 + 1e-8)
+        spec_phase = torch.atan2(imag, real)
+        spec_complex = torch.stack([real, imag], dim=1)[:, :, 1:]  # B,2,256
+        # print("spec", spec_mags.size(), spec_phase.size(), spec_complex.size())
+
+        out = spec_complex
+        encoder_out = []
+        for idx, encoder in enumerate(self.encoder):
+            out = encoder(out)
+            # print("encoder out:", out.size())
+            encoder_out.append(out)
+        B, C, D, T = out.size()
+        out = out.permute(3, 0, 1, 2)
+        if self.use_clstm:
+            r_rnn_in = out[:, :, :C // 2]
+            i_rnn_in = out[:, :, C // 2:]
+            r_rnn_in = torch.reshape(r_rnn_in, [T, B, C // 2 * D])
+            i_rnn_in = torch.reshape(i_rnn_in, [T, B, C // 2 * D])
+
+            r_rnn_in, i_rnn_in = self.enhance([r_rnn_in, i_rnn_in])
+
+            r_rnn_in = torch.reshape(r_rnn_in, [T, B, C // 2, D])
+            i_rnn_in = torch.reshape(i_rnn_in, [T, B, C // 2, D])
+            out = torch.cat([r_rnn_in, i_rnn_in], 2)
+        else:
+            out = torch.reshape(out, [T, B, C * D])
+            out, _ = self.enhance(out)
+            out = self.transform(out)
+            out = torch.reshape(out, [T, B, C, D])
+        out = out.permute(1, 2, 3, 0)
+        for idx in range(len(self.decoder)):
+            out = complex_cat([out, encoder_out[-1 - idx]], 1)
+            out = self.decoder[idx](out)
+            out = out[..., 1:]
+        mask_real = out[:, 0]
+        mask_imag = out[:, 1]
+        mask_real = F.pad(mask_real, [0, 0, 1, 0])
+        mask_imag = F.pad(mask_imag, [0, 0, 1, 0])
+        if self.masking_mode == 'E':
+            mask_mags = (mask_real ** 2 + mask_imag ** 2) ** 0.5
+            real_phase = mask_real / (mask_mags + 1e-8)
+            imag_phase = mask_imag / (mask_mags + 1e-8)
+            mask_phase = torch.atan2(
+                imag_phase,
+                real_phase
+            )
+            mask_mags = torch.tanh(mask_mags)
+            est_mags = mask_mags * spec_mags
+            est_phase = spec_phase + mask_phase
+            real = est_mags * torch.cos(est_phase)
+            imag = est_mags * torch.sin(est_phase)
+        elif self.masking_mode == 'C':
+            real = real * mask_real - imag * mask_imag
+            imag = real * mask_imag + imag * mask_real
+        elif self.masking_mode == 'R':
+            real = real * mask_real
+            imag = imag * mask_imag
+
+        out_spec = torch.cat([real, imag], 1)
+        out_wav = self.istft(out_spec)
+        out_wav = torch.squeeze(out_wav, 1)
+        out_wav = out_wav.clamp_(-1, 1)
+        return out_wav
+
+
+def l2_norm(s1, s2):
+    norm = torch.sum(s1 * s2, -1, keepdim=True)
+    return norm
+
+
+def si_snr(s1, s2, eps=1e-8):
+    s1_s2_norm = l2_norm(s1, s2)
+    s2_s2_norm = l2_norm(s2, s2)
+    s_target = s1_s2_norm / (s2_s2_norm + eps) * s2
+    e_nosie = s1 - s_target
+    target_norm = l2_norm(s_target, s_target)
+    noise_norm = l2_norm(e_nosie, e_nosie)
+    snr = 10 * torch.log10(target_norm / (noise_norm + eps) + eps)
+    return torch.mean(snr)
+
+
+def loss(inputs, label):
+    return -(si_snr(inputs, label))
+
+
+if __name__ == '__main__':
+    test_model = DCCRN(rnn_hidden=256, masking_mode='E', use_clstm=True, kernel_num=(32, 64, 128, 256, 256, 256))
+
+    model_test_timer(test_model, (1, 16000 * 30))