update

examples/clean_unet_aishell/run.sh

@@ -14,8 +14,7 @@ sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name fi
 
 sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
---max_count 10000
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
 
 END

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -76,6 +76,32 @@ 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 = self.relu(x)
+        return x
+
+
 class CleanUNet(nn.Module):
     """
     CleanUNet architecture.
@@ -134,18 +160,39 @@ class CleanUNet(nn.Module):
 
             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)
+                self.decoder.append(DecoderBlock(
+                    channels_h=channels_h,
+                    channels_output=channels_output,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                    do_relu=False,
                 ))
             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(
+                    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)
+            #     ))
             channels_output = channels_h
 
             # double H but keep below max_H