Remove causal option and add batch mask

Author: mthrok

examples/source_separation/conv_tasnet/train.py

@@ -101,7 +101,6 @@ def _get_model(
     msk_num_hidden_feats=512,
     msk_num_layers=8,
     msk_num_stacks=3,
-    causal=False,
 ):
     model = conv_tasnet.model.ConvTasNet(
         num_sources=num_sources,
@@ -112,7 +111,6 @@ def _get_model(
         msk_num_hidden_feats=msk_num_hidden_feats,
         msk_num_layers=msk_num_layers,
         msk_num_stacks=msk_num_stacks,
-        causal=causal,
     )
     _LG.info_on_master("Model Configuration:")
     _LG.info_on_master(" - N: %d", enc_num_feats)

examples/source_separation/conv_tasnet/trainer.py

@@ -98,11 +98,14 @@ def train_one_epoch(self):
 
         num_batches = len(self.train_loader)
         for i, batch in enumerate(self.train_loader, start=1):
-            mixed = batch.mix.to(self.device)
-            sources = batch.src.to(self.device)
+            mix = batch.mix.to(self.device)
+            src = batch.src.to(self.device)
+            mask = batch.mask.to(self.device)
 
-            estimate = self.model(mixed)
-            si_snri, sdri = si_sdr_improvement(estimate, sources, mixed)
+            estimate = self.model(mix)
+            estimate = estimate * mask
+
+            si_snri, sdri = si_sdr_improvement(estimate, src, mix)
             si_snri = si_snri.mean()
             sdri = sdri.mean()
 
@@ -131,29 +134,28 @@ def validate(self):
     def _test(self, loader):
         self.model.eval()
 
-        total_si_snri = 0.0
-        total_sdri = 0.0
+        total_si_snri = torch.zeros(1, dtype=torch.float32, device=self.device)
+        total_sdri = torch.zeros(1, dtype=torch.float32, device=self.device)
 
-        for samples in loader:
-            # Due to the possible length difference, we run evaluation sample-wise
-            for sample in samples:
-                mixed = sample.mix.to(self.device)
-                sources = sample.src.to(self.device)
+        for batch in loader:
+            mix = batch.mix.to(self.device)
+            src = batch.src.to(self.device)
+            mask = batch.mask.to(self.device)
 
-                estimate = self.model(mixed)
-                si_snri, sdri = si_sdr_improvement(estimate, sources, mixed)
-                si_snri = si_snri.sum()
-                sdri = sdri.sum()
+            estimate = self.model(mix)
+            estimate = estimate * mask
 
-                dist.all_reduce(si_snri, dist.ReduceOp.SUM)
-                dist.all_reduce(sdri, dist.ReduceOp.SUM)
+            si_snri, sdri = si_sdr_improvement(estimate, src, mix)
 
-                total_si_snri += si_snri.item()
-                total_sdri += sdri.item()
+            total_si_snri += si_snri.sum()
+            total_sdri += sdri.sum()
 
             if self.debug:
                 break
 
+        dist.all_reduce(total_si_snri, dist.ReduceOp.SUM)
+        dist.all_reduce(total_sdri, dist.ReduceOp.SUM)
+
         num_samples = len(loader.dataset)
-        metric = Metric(total_si_snri / num_samples, total_sdri / num_samples)
+        metric = Metric(total_si_snri.item() / num_samples, total_sdri.item() / num_samples)
         return metric