update

examples/dfnet/step_2_train_model.py

@@ -263,7 +263,7 @@ def main():
 
             est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
 
-            print(f"est_mask.shape: {est_mask.shape}")
+            print(f"est_mask.shape: {est_mask.shape}, est_mask.dtype: {est_mask.dtype}")
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
             mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
             # mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)

toolbox/torchaudio/models/dfnet/modeling_dfnet.py

@@ -892,9 +892,22 @@ class DfNet(nn.Module):
         # spec_e shape: [batch_size, spec_bins, time_steps, 2]
 
         mask = torch.squeeze(mask, dim=1)
-        est_mask = mask.permute(0, 2, 1)
-        # mask shape: [batch_size, spec_bins, time_steps]
+        mask = mask.permute(0, 2, 1)
+        # mask shape: [b, 256, t]
+        est_mask = self.mask_transfer(mask)
+        # est_mask shape: [b, 257, t]
+
+        # spec_e shape: [b, 256, t, 2]
+        est_spec = self.spec_transfer(spec_e)
+        # est_spec shape: [b, 257*2, t]
+        est_wav = self.istft.forward(est_spec)
+        est_wav = torch.squeeze(est_wav, dim=1)
+        est_wav = est_wav[:, :n_samples]
+        # est_wav shape: [b, n_samples]
+        return est_spec, est_wav, est_mask, lsnr
 
+    def spec_transfer(self, spec_e: torch.Tensor) -> torch.Tensor:
+        # spec_e shape: [b, 256, t, 2]
         b, _, t, _ = spec_e.shape
         est_spec = torch.cat(tensors=[
             torch.concat(tensors=[
@@ -906,12 +919,18 @@ class DfNet(nn.Module):
                 torch.zeros(size=(b, 1, t), dtype=spec_e.dtype).to(spec_e.device)
             ], dim=1),
         ], dim=1)
-        # est_spec shape: [b, n+2, t]
-        est_wav = self.istft.forward(est_spec)
-        est_wav = torch.squeeze(est_wav, dim=1)
-        est_wav = est_wav[:, :n_samples]
-        # est_wav shape: [b, n_samples]
-        return est_spec, est_wav, est_mask, lsnr
+        # est_spec shape: [b, 257*2, t]
+        return est_spec
+
+    def mask_transfer(self, mask: torch.Tensor) -> torch.Tensor:
+        # mask shape: [b, 256, t]
+        b, _, t = mask.shape
+        est_mask = torch.concat(tensors=[
+            mask,
+            torch.zeros(size=(b, 1, t), dtype=mask.dtype).to(mask.device)
+        ], dim=1)
+        # est_mask shape: [b, 257, t]
+        return est_mask
 
     def mask_loss_fn(self, est_mask: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
         """
@@ -924,35 +943,18 @@ class DfNet(nn.Module):
         clean_stft = self.stft(clean)
         clean_re = clean_stft[:, :self.freq_bins, :]
         clean_im = clean_stft[:, self.freq_bins:, :]
+        clean_power = clean_re ** 2 + clean_im ** 2
 
         noisy_stft = self.stft(noisy)
         noisy_re = noisy_stft[:, :self.freq_bins, :]
         noisy_im = noisy_stft[:, self.freq_bins:, :]
-
         noisy_power = noisy_re ** 2 + noisy_im ** 2
 
-        sr = clean_re
-        yr = noisy_re
-        si = clean_im
-        yi = noisy_im
-        y_pow = noisy_power
-        # (Sr * Yr + Si * Yi) / (Y_pow + 1e-8)
-        gth_mask_re = (sr * yr + si * yi) / (y_pow + self.eps)
-        # (Si * Yr - Sr * Yi) / (Y_pow + 1e-8)
-        gth_mask_im = (sr * yr - si * yi) / (y_pow + self.eps)
-
-        gth_mask_re[gth_mask_re > 2] = 1
-        gth_mask_re[gth_mask_re < -2] = -1
-        gth_mask_im[gth_mask_im > 2] = 1
-        gth_mask_im[gth_mask_im < -2] = -1
-
-        mask_re = est_mask[:, :self.freq_bins, :]
-        mask_im = est_mask[:, self.freq_bins:, :]
+        speech_irm = clean_power / (noisy_power + self.eps)
+        # speech_irm = torch.pow(speech_irm, self.irm_beta)
 
-        loss_re = F.mse_loss(gth_mask_re, mask_re)
-        loss_im = F.mse_loss(gth_mask_im, mask_im)
+        loss = F.mse_loss(est_mask, speech_irm)
 
-        loss = loss_re + loss_im
         return loss