update

examples/conv_tasnet_gan/step_2_train_model.py

@@ -39,8 +39,8 @@ from tqdm import tqdm
 from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
 from toolbox.torchaudio.models.conv_tasnet.configuration_conv_tasnet import ConvTasNetConfig
 from toolbox.torchaudio.models.conv_tasnet.modeling_conv_tasnet import ConvTasNet, ConvTasNetPretrainedModel
-from toolbox.torchaudio.models.discriminators.conv_tasnet_discriminator.modeling_conv_tasnet_discriminator import ConvTasNetDiscriminatorPretrainedModel
-from toolbox.torchaudio.models.discriminators.conv_tasnet_discriminator.configuration_conv_tasnet_discriminator import ConvTasNetDiscriminatorConfig
+from toolbox.torchaudio.models.discriminators.waveform_metric_discriminator.modeling_waveform_metric_discriminator import WaveformMetricDiscriminatorPretrainedModel
+from toolbox.torchaudio.models.discriminators.waveform_metric_discriminator.configuration_waveform_metric_discriminator import WaveformMetricDiscriminatorConfig
 from toolbox.torchaudio.models.nx_clean_unet.metrics import run_batch_pesq, run_pesq_score
 from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
 from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
@@ -126,7 +126,7 @@ def main():
     config = ConvTasNetConfig.from_pretrained(
         pretrained_model_name_or_path=args.config_file,
     )
-    discriminator_config = ConvTasNetDiscriminatorConfig.from_pretrained(
+    discriminator_config = WaveformMetricDiscriminatorConfig.from_pretrained(
         pretrained_model_name_or_path=args.discriminator_config_file,
     )
 
@@ -189,13 +189,13 @@ def main():
     model.to(device)
     model.train()
 
-    discriminator = ConvTasNetDiscriminatorPretrainedModel(discriminator_config).to(device)
+    discriminator = WaveformMetricDiscriminatorPretrainedModel(discriminator_config).to(device)
     discriminator.to(device)
     discriminator.train()
 
     # optimizer
     logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr, betas=[config.adam_b1, config.adam_b2])
     discriminator_optimizer = torch.optim.AdamW(discriminator.parameters(), config.lr, betas=[config.adam_b1, config.adam_b2])
 
     # resume training

examples/conv_tasnet_gan/yaml/discriminator_config.yaml

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

toolbox/torchaudio/models/discriminators/{conv_tasnet_discriminator/configuration_conv_tasnet_discriminator.py → waveform_metric_discriminator/configuration_waveform_metric_discriminator.py}

@@ -3,13 +3,13 @@
 from toolbox.torchaudio.configuration_utils import PretrainedConfig
 
 
-class ConvTasNetDiscriminatorConfig(PretrainedConfig):
+class WaveformMetricDiscriminatorConfig(PretrainedConfig):
     """
     https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
     """
     def __init__(self,
                  sample_rate: int = 8000,
-                 segment_size: int = 16000,
+                 segment_size: int = 4,
                  n_fft: int = 512,
                  win_length: int = 200,
                  hop_length: int = 80,
@@ -19,7 +19,7 @@ class ConvTasNetDiscriminatorConfig(PretrainedConfig):
 
                  **kwargs
                  ):
-        super(ConvTasNetDiscriminatorConfig, self).__init__(**kwargs)
+        super(WaveformMetricDiscriminatorConfig, self).__init__(**kwargs)
         self.sample_rate = sample_rate
         self.segment_size = segment_size
         self.n_fft = n_fft

toolbox/torchaudio/models/discriminators/{conv_tasnet_discriminator/modeling_conv_tasnet_discriminator.py → waveform_metric_discriminator/modeling_waveform_metric_discriminator.py}

@@ -7,8 +7,7 @@ import torch
 import torch.nn as nn
 import torchaudio
 
-from toolbox.torchaudio.configuration_utils import DISCRIMINATOR_CONFIG_FILE
-from toolbox.torchaudio.models.discriminators.conv_tasnet_discriminator.configuration_conv_tasnet_discriminator import ConvTasNetDiscriminatorConfig
+from toolbox.torchaudio.models.discriminators.waveform_metric_discriminator.configuration_waveform_metric_discriminator import WaveformMetricDiscriminatorConfig
 
 
 class LearnableSigmoid1d(nn.Module):
@@ -23,9 +22,9 @@ class LearnableSigmoid1d(nn.Module):
         return self.beta * torch.sigmoid(self.slope * x)
 
 
-class ConvTasNetDiscriminator(nn.Module):
-    def __init__(self, config: ConvTasNetDiscriminatorConfig):
-        super(ConvTasNetDiscriminator, self).__init__()
+class WaveformMetricDiscriminator(nn.Module):
+    def __init__(self, config: WaveformMetricDiscriminatorConfig):
+        super(WaveformMetricDiscriminator, self).__init__()
         dim = config.discriminator_dim
         self.in_channel = config.discriminator_in_channel
 
@@ -74,21 +73,22 @@ class ConvTasNetDiscriminator(nn.Module):
         return self.layers(xy)
 
 
+CONFIG_FILE = "discriminator_config.yaml"
 MODEL_FILE = "discriminator.pt"
 
 
-class ConvTasNetDiscriminatorPretrainedModel(ConvTasNetDiscriminator):
+class WaveformMetricDiscriminatorPretrainedModel(WaveformMetricDiscriminator):
     def __init__(self,
-                 config: ConvTasNetDiscriminatorConfig,
+                 config: WaveformMetricDiscriminatorConfig,
                  ):
-        super(ConvTasNetDiscriminatorPretrainedModel, self).__init__(
+        super(WaveformMetricDiscriminatorPretrainedModel, self).__init__(
             config=config,
         )
         self.config = config
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-        config = ConvTasNetDiscriminatorConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        config = WaveformMetricDiscriminatorPretrainedModel.from_pretrained(pretrained_model_name_or_path, **kwargs)
 
         model = cls(config)
 
@@ -125,8 +125,8 @@ class ConvTasNetDiscriminatorPretrainedModel(ConvTasNetDiscriminator):
 
 
 def main():
-    config = ConvTasNetDiscriminatorConfig()
-    discriminator = ConvTasNetDiscriminator(config=config)
+    config = WaveformMetricDiscriminatorConfig()
+    discriminator = WaveformMetricDiscriminator(config=config)
 
     # shape: [batch_size, num_samples]
     # x = torch.ones([4, int(4.5 * 16000)])

toolbox/torchaudio/models/discriminators/waveform_metric_discriminator/yaml/discriminator_config.yaml

@@ -0,0 +1,10 @@
+model_name: "waveform_metric_discriminator"
+
+sample_rate: 8000
+segment_size: 4
+n_fft: 512
+win_size: 200
+hop_size: 80
+
+discriminator_dim: 16
+discriminator_in_channel: 2