update

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -76,34 +76,6 @@ def padding(x, D, K, S):
     return x
 
 
-class DecoderBlock(nn.Module):
-    def __init__(self,
-                 channels_h: int,
-                 channels_output: int,
-                 kernel_size: int,
-                 stride: int,
-                 do_relu: bool = True,
-                 ):
-        super(DecoderBlock, self).__init__()
-        self.do_relu = do_relu
-
-        self.conv = nn.Conv1d(channels_h, channels_h * 2, 1)
-        self.glu = nn.GLU(dim=1)
-        self.convt = nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride)
-        # self.relu = nn.ReLU()
-
-    def forward(self, inputs: torch.Tensor):
-        # inputs shape: [batch_size, channel, num_samples]
-        x = self.conv(inputs)
-        x = self.glu(x)
-        x = self.convt(x)
-        if self.do_relu:
-            x = F.relu(x)
-            # x = self.relu(x)
-
-        return x
-
-
 class CleanUNet(nn.Module):
     """
     CleanUNet architecture.
@@ -162,39 +134,18 @@ class CleanUNet(nn.Module):
 
             if i == 0:
                 # no relu at end
-                self.decoder.append(DecoderBlock(
-                    channels_h=channels_h,
-                    channels_output=channels_output,
-                    kernel_size=kernel_size,
-                    stride=stride,
-                    do_relu=False,
+                self.decoder.append(nn.Sequential(
+                    nn.Conv1d(channels_h, channels_h * 2, 1),
+                    nn.GLU(dim=1),
+                    nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride)
                 ))
             else:
-                self.decoder.insert(
-                    index=0,
-                    module=DecoderBlock(
-                        channels_h=channels_h,
-                        channels_output=channels_output,
-                        kernel_size=kernel_size,
-                        stride=stride,
-                        do_relu=True,
-                    )
-                )
-
-            # if i == 0:
-            #     # no relu at end
-            #     self.decoder.append(nn.Sequential(
-            #         nn.Conv1d(channels_h, channels_h * 2, 1),
-            #         nn.GLU(dim=1),
-            #         nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride)
-            #     ))
-            # else:
-            #     self.decoder.insert(0, nn.Sequential(
-            #         nn.Conv1d(channels_h, channels_h * 2, 1),
-            #         nn.GLU(dim=1),
-            #         nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride),
-            #         # nn.ReLU(inplace=False)
-            #     ))
+                self.decoder.insert(0, nn.Sequential(
+                    nn.Conv1d(channels_h, channels_h * 2, 1),
+                    nn.GLU(dim=1),
+                    nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride),
+                    nn.ReLU()
+                ))
             channels_output = channels_h
 
             # double H but keep below max_H