Add target cropping if outputs are different length

remfx/models.py

@@ -13,6 +13,7 @@ from remfx.utils import FADLoss, spectrogram
 from remfx.dptnet import DPTNet_base
 from remfx.dcunet import RefineSpectrogramUnet
 from remfx.tcn import TCN
+from remfx.utils import causal_crop
 
 
 class RemFX(pl.LightningModule):
@@ -223,21 +224,14 @@ class DCUNetModel(nn.Module):
     def forward(self, batch):
         x, target = batch
         output = self.model(x.squeeze(1)).unsqueeze(1)  # B x 1 x T
-        # Pad or crop to match target
-        if output.shape[-1] > target.shape[-1]:
-            output = output[:, : target.shape[-1]]
-        elif output.shape[-1] < target.shape[-1]:
-            output = F.pad(output, (0, target.shape[-1] - output.shape[-1]))
+        # Crop target to match output
+        if output.shape[-1] < target.shape[-1]:
+            target = causal_crop(target, output.shape[-1])
         loss = self.mrstftloss(output, target) + self.l1loss(output, target) * 100
         return loss, output
 
     def sample(self, x: Tensor) -> Tensor:
         output = self.model(x.squeeze(1)).unsqueeze(1)  # B x 1 x T
-        # Pad or crop to match target
-        if output.shape[-1] > x.shape[-1]:
-            output = output[:, : x.shape[-1]]
-        elif output.shape[-1] < x.shape[-1]:
-            output = F.pad(output, (0, x.shape[-1] - output.shape[-1]))
         return output
 
 
@@ -253,21 +247,14 @@ class TCNModel(nn.Module):
     def forward(self, batch):
         x, target = batch
         output = self.model(x)  # B x 1 x T
-        # Pad or crop to match target
-        if output.shape[-1] > x.shape[-1]:
-            output = output[:, : x.shape[-1]]
-        elif output.shape[-1] < x.shape[-1]:
-            output = F.pad(output, (0, x.shape[-1] - output.shape[-1]))
+        # Crop target to match output
+        if output.shape[-1] < target.shape[-1]:
+            target = causal_crop(target, output.shape[-1])
         loss = self.mrstftloss(output, target) + self.l1loss(output, target) * 100
         return loss, output
 
     def sample(self, x: Tensor) -> Tensor:
         output = self.model(x)  # B x 1 x T
-        # Pad or crop to match target
-        if output.shape[-1] > x.shape[-1]:
-            output = output[:, : x.shape[-1]]
-        elif output.shape[-1] < x.shape[-1]:
-            output = F.pad(output, (0, x.shape[-1] - output.shape[-1]))
         return output
 
 

remfx/tcn.py

@@ -25,8 +25,6 @@ class TCNBlock(nn.Module):
         self.stride = stride
 
         self.crop_fn = crop_fn
-        # Assumes stride of 1
-        padding = (kernel_size + (kernel_size - 1) * (dilation - 1) - 1) // 2
         self.conv1 = nn.Conv1d(
             in_ch,
             out_ch,