update

examples/mpnet_aishell/step_2_train_model.py

@@ -32,7 +32,7 @@ from tqdm import tqdm
 
 from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
 from toolbox.torchaudio.models.mpnet.configuation_mpnet import MPNetConfig
-from toolbox.torchaudio.models.mpnet.discriminator import MetricDiscriminator, batch_pesq
+from toolbox.torchaudio.models.mpnet.discriminator import MetricDiscriminatorPretrainedModel, batch_pesq
 from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNet, MPNetPretrainedModel, phase_losses, pesq_score
 from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft
 
@@ -164,14 +164,14 @@ def main():
     # models
     logger.info(f"prepare models. config_file: {args.config_file}")
     generator = MPNetPretrainedModel(config).to(device)
-    discriminator = MetricDiscriminator().to(device)
+    discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
 
     # optimizer
-    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    logger.info("prepare optimizer, lr_scheduler")
     num_params = 0
     for p in generator.parameters():
         num_params += p.numel()
-    print("Total Parameters (generator): {:.3f}M".format(num_params/1e6))
+    logger.info("Total Parameters (generator): {:.3f}M".format(num_params/1e6))
 
     optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
     optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
@@ -180,8 +180,24 @@ def main():
     scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=-1)
 
     # training loop
+
+    # state
+    loss_d = 10000000000
+    loss_g = 10000000000
+    pesq_metric = 10000000000
+    mag_err = 10000000000
+    pha_err = 10000000000
+    com_err = 10000000000
+    stft_err = 10000000000
+
+    model_list = list()
+    best_idx_epoch = None
+    best_metric = None
+    patience_count = 0
+
     logger.info("training")
     for idx_epoch in range(args.max_epochs):
+        # train
         generator.train()
         discriminator.train()
 
@@ -251,12 +267,16 @@ def main():
             total_loss_g += loss_gen_all.item()
             total_batches += 1
 
+            loss_d = round(total_loss_d / total_batches, 4)
+            loss_g = round(total_loss_g / total_batches, 4)
+
             progress_bar.update(1)
             progress_bar.set_postfix({
-                "loss_d": round(total_loss_d / total_batches, 4),
-                "loss_g": round(total_loss_g / total_batches, 4),
+                "loss_d": loss_d,
+                "loss_g": loss_g,
             })
 
+        # evaluation
         generator.eval()
         torch.cuda.empty_cache()
         total_pesq_score = 0.
@@ -297,18 +317,87 @@ def main():
 
                 total_batches += 1
 
+                pesq_metric = round(total_pesq_score / total_batches, 4)
+                mag_err = round(total_mag_err / total_batches, 4)
+                pha_err = round(total_pha_err / total_batches, 4)
+                com_err = round(total_com_err / total_batches, 4)
+                stft_err = round(total_stft_err / total_batches, 4)
+
                 progress_bar.update(1)
                 progress_bar.set_postfix({
-                    "pesq_score": round(total_pesq_score / total_batches, 4),
-                    "mag_err": round(total_mag_err / total_batches, 4),
-                    "pha_err": round(total_pha_err / total_batches, 4),
-                    "com_err": round(total_com_err / total_batches, 4),
-                    "stft_err": round(total_stft_err / total_batches, 4),
-
+                    "pesq_metric": pesq_metric,
+                    "mag_err": mag_err,
+                    "pha_err": pha_err,
+                    "com_err": com_err,
+                    "stft_err": stft_err,
                 })
 
+        # scheduler
+        scheduler_g.step()
+        scheduler_d.step()
+
+        # save path
+        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+        epoch_dir.mkdir(parents=True, exist_ok=False)
+
+        # save models
+        generator.save_pretrained(epoch_dir.as_posix())
+        discriminator.save_pretrained(epoch_dir.as_posix())
+
+        model_list.append(epoch_dir)
+        if len(model_list) >= args.num_serialized_models_to_keep:
+            model_to_delete: Path = model_list.pop(0)
+            shutil.rmtree(model_to_delete.as_posix())
+
+        # save metric
+        if best_metric is None:
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        elif pesq_metric < best_metric:
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        else:
+            pass
+
+        metrics = {
+            "idx_epoch": idx_epoch,
+            "best_idx_epoch": best_idx_epoch,
+            "loss_d": loss_d,
+            "loss_g": loss_g,
+
+            "pesq_metric": pesq_metric,
+            "mag_err": mag_err,
+            "pha_err": pha_err,
+            "com_err": com_err,
+            "stft_err": stft_err,
+
+        }
+        metrics_filename = epoch_dir / "metrics_epoch.json"
+        with open(metrics_filename, "w", encoding="utf-8") as f:
+            json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+        # save best
+        best_dir = serialization_dir / "best"
+        if best_idx_epoch == idx_epoch:
+            if best_dir.exists():
+                shutil.rmtree(best_dir)
+            shutil.copytree(epoch_dir, best_dir)
+
+        # early stop
+        early_stop_flag = False
+        if best_idx_epoch == idx_epoch:
+            patience_count = 0
+        else:
+            patience_count += 1
+        if patience_count >= args.patience:
+            early_stop_flag = True
+
+        # early stop
+        if early_stop_flag:
+            break
+
     return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

toolbox/torchaudio/models/mpnet/configuation_mpnet.py

@@ -33,6 +33,9 @@ class MPNetConfig(PretrainedConfig):
 
                  dist_config: dict = None,
 
+                 discriminator_dim: int = 32,
+                 discriminator_in_channel: int = 2,
+
                  **kwargs
                  ):
         super(MPNetConfig, self).__init__(**kwargs)
@@ -63,6 +66,9 @@ class MPNetConfig(PretrainedConfig):
             "world_size": 1
         }
 
+        self.discriminator_dim = discriminator_dim
+        self.discriminator_in_channel = discriminator_in_channel
+
 
 if __name__ == "__main__":
     pass

toolbox/torchaudio/models/mpnet/discriminator.py

@@ -1,5 +1,8 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
+import os
+from typing import Optional, Union
+
 import torch
 import torch.nn as nn
 import numpy as np
@@ -7,6 +10,8 @@ import torch.nn.functional as F
 from pesq import pesq
 from joblib import Parallel, delayed
 
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.mpnet.configuation_mpnet import MPNetConfig
 from toolbox.torchaudio.models.mpnet.utils import LearnableSigmoid1d
 
 
@@ -38,8 +43,12 @@ def metric_loss(metric_ref, metrics_gen):
 
 
 class MetricDiscriminator(nn.Module):
-    def __init__(self, dim=16, in_channel=2):
+    def __init__(self, config: MPNetConfig):
         super(MetricDiscriminator, self).__init__()
+
+        dim = config.discriminator_dim
+        in_channel = config.discriminator_in_channel
+
         self.layers = nn.Sequential(
             nn.utils.spectral_norm(nn.Conv2d(in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
             nn.InstanceNorm2d(dim, affine=True),
@@ -67,5 +76,54 @@ class MetricDiscriminator(nn.Module):
         return self.layers(xy)
 
 
+MODEL_FILE = "discriminator.pt"
+
+
+class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
+    def __init__(self,
+                 config: MPNetConfig,
+                 ):
+        super(MetricDiscriminatorPretrainedModel, self).__init__(
+            config=config,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = MPNetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
 if __name__ == '__main__':
     pass

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -183,7 +183,7 @@ class MPNet(nn.Module):
         return denoised_amp, denoised_pha, denoised_com
 
 
-MODEL_FILE = "model.pt"
+MODEL_FILE = "generator.pt"
 
 
 class MPNetPretrainedModel(MPNet):