Spaces:
Runtime error
Runtime error
| # 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)) | |