update

examples/conv_tasnet_gan/step_2_train_model.py

@@ -211,9 +211,9 @@ def main():
     if last_step_idx != -1:
         logger.info(f"resume from steps-{last_step_idx}.")
         model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
-        discriminator_pt = serialization_dir / f"steps-{last_step_idx}/discriminator.pt"
-
         optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
+
+        discriminator_pt = serialization_dir / f"steps-{last_step_idx}/discriminator.pt"
         discriminator_optimizer_pth = serialization_dir / f"steps-{last_step_idx}/discriminator_optimizer.pth"
 
         logger.info(f"load state dict for model.")
@@ -221,10 +221,11 @@ def main():
             state_dict = torch.load(f, map_location="cpu", weights_only=True)
         model.load_state_dict(state_dict, strict=True)
 
-        logger.info(f"load state dict for optimizer.")
-        with open(optimizer_pth.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        optimizer.load_state_dict(state_dict)
+        if optimizer_pth.exists():
+            logger.info(f"load state dict for optimizer.")
+            with open(optimizer_pth.as_posix(), "rb") as f:
+                state_dict = torch.load(f, map_location="cpu", weights_only=True)
+            optimizer.load_state_dict(state_dict)
 
         if discriminator_pt.exists():
             logger.info(f"load state dict for discriminator.")
@@ -497,6 +498,8 @@ def main():
                 total_neg_stoi_loss = 0.
                 total_mr_stft_loss = 0.
                 total_pesq_loss = 0.
+                total_discriminator_g_loss = 0.
+                total_discriminator_d_loss = 0.
                 total_batches = 0.
 
                 progress_bar_eval.close()

examples/conv_tasnet_gan/yaml/config.yaml

@@ -5,11 +5,11 @@ segment_size: 4
 
 win_size: 20
 freq_bins: 256
-bottleneck_channels: 256
+bottleneck_channels: 128
 num_speakers: 1
-num_blocks: 4
-num_sub_blocks: 8
-sub_blocks_channels: 512
+num_blocks: 2
+num_sub_blocks: 4
+sub_blocks_channels: 256
 sub_blocks_kernel_size: 3
 
 norm_type: "gLN"

examples/conv_tasnet_gan/yaml/discriminator_config.yaml

@@ -6,5 +6,5 @@ n_fft: 512
 win_size: 200
 hop_size: 80
 
-discriminator_dim: 24
+discriminator_dim: 32
 discriminator_in_channel: 2