update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -21,6 +21,7 @@ sys.path.append(os.path.join(pwd, "../../"))
 import numpy as np
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 from torch.utils.data.dataloader import DataLoader
 import torchaudio
 from tqdm import tqdm
@@ -95,6 +96,28 @@ class CollateFunction(object):
             window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
         )
 
+    @staticmethod
+    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
+        batch_size, channels, freq_dim, time_steps = x.shape
+
+        # kernel: [freq_dim, n_time_step]
+        kernel_size = (freq_dim, n_time_steps)
+
+        # pad
+        pad = n_time_steps // 2
+        x = torch.concat(tensors=[
+            x[:, :, :, :pad],
+            x,
+            x[:, :, :, -pad:],
+        ], dim=-1)
+
+        x = F.unfold(
+            input=x,
+            kernel_size=kernel_size,
+        )
+        # x shape: [batch_size, fold, time_steps]
+        return x
+
     def __call__(self, batch: List[dict]):
         mix_spec_list = list()
         speech_irm_list = list()
@@ -119,7 +142,13 @@ class CollateFunction(object):
             snr_db: torch.Tensor = 10 * torch.log10(
                 speech_spec / (noise_spec + self.epsilon)
             )
-            snr_db = torch.mean(snr_db, dim=0, keepdim=True)
+            snr_db_ = torch.unsqueeze(snr_db, dim=0)
+            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
+            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
+            snr_db_ = torch.squeeze(snr_db_, dim=0)
+            # snr_db_ shape: [fold, time_steps]
+
+            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
             # snr_db shape: [1, time_steps]
 
             mix_spec_list.append(mix_spec)
@@ -262,9 +291,9 @@ def main():
 
             speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-            # loss = irm_loss + 0.1 * snr_loss
-            loss = irm_loss
+            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+            loss = irm_loss + 0.1 * snr_loss
+            # loss = irm_loss
 
             total_loss += loss.item()
             total_examples += mix_spec.size(0)
@@ -297,9 +326,9 @@ def main():
             with torch.no_grad():
                 speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
                 irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-                # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-                # loss = irm_loss + 0.1 * snr_loss
-                loss = irm_loss
+                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+                loss = irm_loss + 0.1 * snr_loss
+                # loss = irm_loss
 
             total_loss += loss.item()
             total_examples += mix_spec.size(0)