update

examples/clean_unet_aishell/step_2_train_model.py

@@ -2,6 +2,8 @@
 # -*- coding: utf-8 -*-
 """
 https://github.com/NVIDIA/CleanUNet/blob/main/train.py
+
+https://github.com/NVIDIA/CleanUNet/blob/main/configs/DNS-large-full.json
 """
 import argparse
 import json
@@ -20,6 +22,7 @@ sys.path.append(os.path.join(pwd, "../../"))
 
 import numpy as np
 import torch
+import torch.nn as nn
 from torch.nn import functional as F
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
@@ -27,6 +30,9 @@ from tqdm import tqdm
 from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
 from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUnetConfig
 from toolbox.torchaudio.models.clean_unet.modeling_clean_unet import CleanUNetPretrainedModel
+from toolbox.torchaudio.models.clean_unet.training import LinearWarmupCosineDecay
+from toolbox.torchaudio.models.clean_unet.loss import MultiResolutionSTFTLoss
+from toolbox.torchaudio.models.clean_unet.metrics import batch_pesq
 
 
 def get_args():
@@ -36,6 +42,9 @@ def get_args():
 
     parser.add_argument("--max_epochs", default=100, type=int)
 
+    parser.add_argument("--batch_size", default=64, type=int)
+    parser.add_argument("--learning_rate", default=2e-4, type=float)
+
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
@@ -158,56 +167,37 @@ def main():
     model = CleanUNetPretrainedModel(config).to(device)
 
     # optimizer
-    logger.info("prepare optimizer, lr_scheduler")
-    optim_g = torch.optim.AdamW(model.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
-
-    # resume training
-    last_epoch = -1
-    for epoch_i in serialization_dir.glob("epoch-*"):
-        epoch_i = Path(epoch_i)
-        epoch_idx = epoch_i.stem.split("-")[1]
-        epoch_idx = int(epoch_idx)
-        if epoch_idx > last_epoch:
-            last_epoch = epoch_idx
-
-    if last_epoch != -1:
-        logger.info(f"resume from epoch-{last_epoch}.")
-        generator_pt = serialization_dir / f"epoch-{last_epoch}/generator.pt"
-        discriminator_pt = serialization_dir / f"epoch-{last_epoch}/discriminator.pt"
-        optim_g_pth = serialization_dir / f"epoch-{last_epoch}/optim_g.pth"
-        optim_d_pth = serialization_dir / f"epoch-{last_epoch}/optim_d.pth"
-
-        logger.info(f"load state dict for generator.")
-        with open(generator_pt.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        generator.load_state_dict(state_dict, strict=True)
-        logger.info(f"load state dict for discriminator.")
-        with open(discriminator_pt.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        discriminator.load_state_dict(state_dict, strict=True)
-
-        logger.info(f"load state dict for optim_g.")
-        with open(optim_g_pth.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        optim_g.load_state_dict(state_dict)
-        logger.info(f"load state dict for optim_d.")
-        with open(optim_d_pth.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        optim_d.load_state_dict(state_dict)
-
-    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
-    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    optimizer = torch.optim.AdamW(model.parameters(), config.learning_rate)
+    lr_scheduler = LinearWarmupCosineDecay(
+        optimizer,
+        lr_max=args.learning_rate,
+        n_iter=250000,
+        iteration=250000,
+        divider=25,
+        warmup_proportion=0.05,
+        phase=("linear", "cosine"),
+    )
+    # ae_loss_fn = nn.MSELoss(reduction="mean")
+    ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
+
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_sizes=[512, 1024, 2048],
+        hop_sizes=[50, 120, 240],
+        win_lengths=[240, 600, 1200],
+        sc_lambda=0.5,
+        mag_lambda=0.5,
+        band="full"
+    ).to(device)
 
     # training loop
 
     # state
-    loss_d = 10000000000
-    loss_g = 10000000000
-    pesq_metric = 10000000000
-    mag_err = 10000000000
-    pha_err = 10000000000
-    com_err = 10000000000
-    stft_err = 10000000000
+    average_pesq_metric = 10000000000
+    average_loss = 10000000000
+    average_ae_loss = 10000000000
+    average_sc_loss = 10000000000
+    average_mag_loss = 10000000000
 
     model_list = list()
     best_idx_epoch = None
@@ -215,96 +205,74 @@ def main():
     patience_count = 0
 
     logger.info("training")
-    for idx_epoch in range(max(0, last_epoch+1), args.max_epochs):
+    for idx_epoch in range(args.max_epochs):
         # train
-        generator.train()
-        discriminator.train()
+        model.train()
 
-        total_loss_d = 0.
-        total_loss_g = 0.
+        total_pesq_metric = 0.
+        total_loss = 0.
+        total_ae_loss = 0.
+        total_sc_loss = 0.
+        total_mag_loss = 0.
         total_batches = 0.
+
         progress_bar = tqdm(
             total=len(train_data_loader),
             desc="Training; epoch: {}".format(idx_epoch),
         )
         for batch in train_data_loader:
-            clean_audio, noisy_audio = batch
-            clean_audio = clean_audio.to(device)
-            noisy_audio = noisy_audio.to(device)
-            one_labels = torch.ones(clean_audio.shape[0]).to(device)
-
-            clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-            noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-
-            mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
-
-            audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-            mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-
-            audio_list_r, audio_list_g = list(clean_audio.cpu().numpy()), list(audio_g.detach().cpu().numpy())
-            batch_pesq_score = batch_pesq(audio_list_r, audio_list_g)
-
-            # Discriminator
-            optim_d.zero_grad()
-            metric_r = discriminator.forward(clean_mag, clean_mag)
-            metric_g = discriminator.forward(clean_mag, mag_g_hat.detach())
-            loss_disc_r = F.mse_loss(one_labels, metric_r.flatten())
-
-            if batch_pesq_score is not None:
-                loss_disc_g = F.mse_loss(batch_pesq_score.to(device), metric_g.flatten())
-            else:
-                # print("pesq is None!")
-                loss_disc_g = 0
-
-            loss_disc_all = loss_disc_r + loss_disc_g
-            loss_disc_all.backward()
-            optim_d.step()
-
-            # Generator
-            optim_g.zero_grad()
-            # L2 Magnitude Loss
-            loss_mag = F.mse_loss(clean_mag, mag_g)
-            # Anti-wrapping Phase Loss
-            loss_ip, loss_gd, loss_iaf = phase_losses(clean_pha, pha_g)
-            loss_pha = loss_ip + loss_gd + loss_iaf
-            # L2 Complex Loss
-            loss_com = F.mse_loss(clean_com, com_g) * 2
-            # L2 Consistency Loss
-            loss_stft = F.mse_loss(com_g, com_g_hat) * 2
-            # Time Loss
-            loss_time = F.l1_loss(clean_audio, audio_g)
-            # Metric Loss
-            metric_g = discriminator.forward(clean_mag, mag_g_hat)
-            loss_metric = F.mse_loss(metric_g.flatten(), one_labels)
-
-            loss_gen_all = loss_mag * 0.9 + loss_pha * 0.3  + loss_com * 0.1 + loss_stft * 0.1 + loss_metric * 0.05 + loss_time * 0.2
-
-            loss_gen_all.backward()
-            optim_g.step()
-
-            total_loss_d += loss_disc_all.item()
-            total_loss_g += loss_gen_all.item()
+            clean_audios, noisy_audios = batch
+            clean_audios = clean_audios.to(device)
+            noisy_audios = noisy_audios.to(device)
+
+            enhanced_audios = model.forward(noisy_audios)
+
+            ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
+            sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios.squeeze(1), clean_audios.squeeze(1))
+
+            loss = ae_loss + sc_loss + mag_loss
+
+            enhanced_audios_list_r = list(enhanced_audios.cpu().numpy())
+            clean_audios_list_r = list(clean_audios.cpu().numpy())
+            pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_metric += pesq_metric.item()
+            total_loss += loss.item()
+            total_ae_loss += ae_loss.item()
+            total_sc_loss += sc_loss.item()
+            total_mag_loss += mag_loss.item()
             total_batches += 1
 
-            loss_d = round(total_loss_d / total_batches, 4)
-            loss_g = round(total_loss_g / total_batches, 4)
+            average_pesq_metric = round(total_pesq_metric / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_ae_loss = round(total_ae_loss / total_batches, 4)
+            average_sc_loss = round(total_sc_loss / total_batches, 4)
+            average_mag_loss = round(total_mag_loss / total_batches, 4)
 
             progress_bar.update(1)
             progress_bar.set_postfix({
-                "loss_d": loss_d,
-                "loss_g": loss_g,
+                "pesq_metric": average_pesq_metric,
+                "loss": average_loss,
+                "ae_loss": average_ae_loss,
+                "sc_loss": average_sc_loss,
+                "mag_loss": average_mag_loss,
             })
 
         # evaluation
-        generator.eval()
-        discriminator.eval()
+        model.eval()
 
         torch.cuda.empty_cache()
-        total_pesq_score = 0.
-        total_mag_err = 0.
-        total_pha_err = 0.
-        total_com_err = 0.
-        total_stft_err = 0.
+
+        total_pesq_metric = 0.
+        total_loss = 0.
+        total_ae_loss = 0.
+        total_sc_loss = 0.
+        total_mag_loss = 0.
         total_batches = 0.
 
         progress_bar = tqdm(
@@ -313,61 +281,52 @@ def main():
         )
         with torch.no_grad():
             for batch in valid_data_loader:
-                clean_audio, noisy_audio = batch
-                clean_audio = clean_audio.to(device)
-                noisy_audio = noisy_audio.to(device)
-
-                clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-                noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-
-                mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
-
-                audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-                mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
-
-                total_pesq_score += pesq_score(
-                    torch.split(clean_audio, 1, dim=0),
-                    torch.split(audio_g, 1, dim=0),
-                    config
-                ).item()
-                total_mag_err += F.mse_loss(clean_mag, mag_g).item()
-                val_ip_err, val_gd_err, val_iaf_err = phase_losses(clean_pha, pha_g)
-                total_pha_err += (val_ip_err + val_gd_err + val_iaf_err).item()
-                total_com_err += F.mse_loss(clean_com, com_g).item()
-                total_stft_err += F.mse_loss(com_g, com_g_hat).item()
-
+                clean_audios, noisy_audios = batch
+                clean_audios = clean_audios.to(device)
+                noisy_audios = noisy_audios.to(device)
+
+                enhanced_audios = model.forward(noisy_audios)
+                enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+                ae_loss = ae_loss_fn(enhanced_audios, enhanced_audios)
+                sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios.squeeze(1), clean_audios.squeeze(1))
+
+                loss = ae_loss + sc_loss + mag_loss
+
+                enhanced_audios_list_r = list(enhanced_audios.cpu().numpy())
+                clean_audios_list_r = list(clean_audios.cpu().numpy())
+                pesq_metric = batch_pesq(enhanced_audios_list_r, clean_audios_list_r)
+
+                total_pesq_metric += pesq_metric.item()
+                total_loss += loss.item()
+                total_ae_loss += ae_loss.item()
+                total_sc_loss += sc_loss.item()
+                total_mag_loss += mag_loss.item()
                 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)
+                average_pesq_metric = round(total_pesq_metric / total_batches, 4)
+                average_loss = round(total_loss / total_batches, 4)
+                average_ae_loss = round(total_ae_loss / total_batches, 4)
+                average_sc_loss = round(total_sc_loss / total_batches, 4)
+                average_mag_loss = round(total_mag_loss / total_batches, 4)
 
                 progress_bar.update(1)
                 progress_bar.set_postfix({
-                    "pesq_metric": pesq_metric,
-                    "mag_err": mag_err,
-                    "pha_err": pha_err,
-                    "com_err": com_err,
-                    "stft_err": stft_err,
+                    "pesq_metric": average_pesq_metric,
+                    "loss": average_loss,
+                    "ae_loss": average_ae_loss,
+                    "sc_loss": average_sc_loss,
+                    "mag_loss": average_mag_loss,
                 })
 
         # scheduler
-        scheduler_g.step()
-        scheduler_d.step()
+        lr_scheduler.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())
-
-        # save optim
-        torch.save(optim_d.state_dict(), (epoch_dir / "optim_d.pth").as_posix())
-        torch.save(optim_g.state_dict(), (epoch_dir / "optim_g.pth").as_posix())
+        model.save_pretrained(epoch_dir.as_posix())
 
         model_list.append(epoch_dir)
         if len(model_list) >= args.num_serialized_models_to_keep:
@@ -377,25 +336,23 @@ def main():
         # save metric
         if best_metric is None:
             best_idx_epoch = idx_epoch
-            best_metric = pesq_metric
-        elif pesq_metric > best_metric:
+            best_metric = average_pesq_metric
+        elif average_pesq_metric > best_metric:
             # great is better.
             best_idx_epoch = idx_epoch
-            best_metric = pesq_metric
+            best_metric = average_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,
+
+            "pesq_metric": average_pesq_metric,
+            "loss": average_loss,
+            "ae_loss": average_ae_loss,
+            "sc_loss": average_sc_loss,
+            "mag_loss": average_mag_loss,
 
         }
         metrics_filename = epoch_dir / "metrics_epoch.json"

toolbox/torchaudio/models/clean_unet/loss.py

@@ -0,0 +1,171 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+import torch
+import torch.nn.functional as F
+
+# from distutils.version import LooseVersion
+
+
+# is_pytorch_17plus = LooseVersion(torch.__version__) >= LooseVersion("1.7")
+is_pytorch_17plus = True
+
+
+def stft(x, fft_size, hop_size, win_length, window):
+    """
+    Perform STFT and convert to magnitude spectrogram.
+    :param x: Tensor, Input signal tensor (B, T).
+    :param fft_size: int, FFT size.
+    :param hop_size: int, Hop size.
+    :param win_length: int, Window length.
+    :param window: str, Window function type.
+    :return: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
+    """
+
+    if is_pytorch_17plus:
+        x_stft = torch.stft(
+            x, fft_size, hop_size, win_length, window, return_complex=False
+        )
+    else:
+        x_stft = torch.stft(x, fft_size, hop_size, win_length, window)
+    real = x_stft[..., 0]
+    imag = x_stft[..., 1]
+
+    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
+    return torch.sqrt(torch.clamp(real**2 + imag**2, min=1e-7)).transpose(2, 1)
+
+
+class SpectralConvergenceLoss(torch.nn.Module):
+    """Spectral convergence loss module."""
+
+    def __init__(self):
+        super(SpectralConvergenceLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """
+        Calculate forward propagation.
+        :param x_mag: Tensor, Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+        :param y_mag: Tensor, Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
+        :return: Tensor, Spectral convergence loss value.
+        """
+        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
+
+
+class LogSTFTMagnitudeLoss(torch.nn.Module):
+    """Log STFT magnitude loss module."""
+
+    def __init__(self):
+        super(LogSTFTMagnitudeLoss, self).__init__()
+
+    def forward(self, x_mag, y_mag):
+        """
+        Calculate forward propagation.
+        :param x_mag: Tensor, Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
+        :param y_mag: Tensor, Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
+        :return: Tensor, Log STFT magnitude loss value.
+        """
+        return F.l1_loss(torch.log(y_mag), torch.log(x_mag))
+
+
+class STFTLoss(torch.nn.Module):
+    """STFT loss module."""
+
+    def __init__(
+        self, fft_size=1024, shift_size=120, win_length=600, window="hann_window",
+        band="full"
+    ):
+        super(STFTLoss, self).__init__()
+        self.fft_size = fft_size
+        self.shift_size = shift_size
+        self.win_length = win_length
+        self.band = band
+
+        self.spectral_convergence_loss = SpectralConvergenceLoss()
+        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss()
+        # NOTE(kan-bayashi): Use register_buffer to fix #223
+        self.register_buffer("window", getattr(torch, window)(win_length))
+
+    def forward(self, x, y):
+        """
+        Calculate forward propagation.
+        :param x: Tensor, Predicted signal (B, T).
+        :param y: Tensor, Groundtruth signal (B, T).
+        :return:
+        Tensor, Spectral convergence loss value.
+        Tensor, Log STFT magnitude loss value.
+        """
+        x_mag = stft(x, self.fft_size, self.shift_size, self.win_length, self.window)
+        y_mag = stft(y, self.fft_size, self.shift_size, self.win_length, self.window)
+
+        if self.band == "high":
+            freq_mask_ind = x_mag.shape[1] // 2  # only select high frequency bands
+            sc_loss  = self.spectral_convergence_loss(x_mag[:,freq_mask_ind:,:], y_mag[:,freq_mask_ind:,:])
+            mag_loss = self.log_stft_magnitude_loss(x_mag[:,freq_mask_ind:,:], y_mag[:,freq_mask_ind:,:])
+        elif self.band == "full":
+            sc_loss  = self.spectral_convergence_loss(x_mag, y_mag)
+            mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)
+        else:
+            raise NotImplementedError
+
+        return sc_loss, mag_loss
+
+
+class MultiResolutionSTFTLoss(torch.nn.Module):
+    """Multi resolution STFT loss module."""
+
+    def __init__(self,
+                 fft_sizes=None, hop_sizes=None, win_lengths=None,
+                 window="hann_window", sc_lambda=0.1, mag_lambda=0.1, band="full",
+                 ):
+        """
+        Initialize Multi resolution STFT loss module.
+        :param fft_sizes: list, List of FFT sizes.
+        :param hop_sizes: list, List of hop sizes.
+        :param win_lengths: list, List of window lengths.
+        :param window: str, Window function type.
+        :param sc_lambda: float, a balancing factor across different losses.
+        :param mag_lambda: float, a balancing factor across different losses.
+        :param band: str, high-band or full-band loss
+        """
+        super(MultiResolutionSTFTLoss, self).__init__()
+        fft_sizes = fft_sizes or [1024, 2048, 512]
+        hop_sizes = hop_sizes or [120, 240, 50]
+        win_lengths = win_lengths or [600, 1200, 240]
+
+        self.sc_lambda = sc_lambda
+        self.mag_lambda = mag_lambda
+
+        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
+        self.stft_losses = torch.nn.ModuleList()
+        for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths):
+            self.stft_losses += [STFTLoss(fs, ss, wl, window, band)]
+
+    def forward(self, x, y):
+        """
+        Calculate forward propagation.
+        :param x: Tensor, Predicted signal (B, T) or (B, #subband, T).
+        :param y: Tensor, Groundtruth signal (B, T) or (B, #subband, T).
+        :return:
+        Tensor, Multi resolution spectral convergence loss value.
+        Tensor, Multi resolution log STFT magnitude loss value.
+        """
+        if len(x.shape) == 3:
+            x = x.view(-1, x.size(2))  # (B, C, T) -> (B x C, T)
+            y = y.view(-1, y.size(2))  # (B, C, T) -> (B x C, T)
+        sc_loss = 0.0
+        mag_loss = 0.0
+        for f in self.stft_losses:
+            sc_l, mag_l = f(x, y)
+            sc_loss += sc_l
+            mag_loss += mag_l
+
+        sc_loss *= self.sc_lambda
+        sc_loss /= len(self.stft_losses)
+        mag_loss *= self.mag_lambda
+        mag_loss /= len(self.stft_losses)
+
+        return sc_loss, mag_loss
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/clean_unet/metrics.py

@@ -0,0 +1,42 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from joblib import Parallel, delayed
+import numpy as np
+from pesq import pesq
+import torch
+
+
+def cal_pesq(clean, noisy, sr=16000):
+    try:
+        pesq_score = pesq(sr, clean, noisy, "wb")
+    except Exception as e:
+        # print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
+        # error can happen due to silent period
+        pesq_score = -1
+    return pesq_score
+
+
+def batch_pesq(clean, noisy):
+    pesq_score = Parallel(n_jobs=15)(delayed(cal_pesq)(c, n) for c, n in zip(clean, noisy))
+    pesq_score = np.array(pesq_score)
+    if -1 in pesq_score:
+        return None
+    pesq_score = (pesq_score - 1) / 3.5
+    return torch.FloatTensor(pesq_score)
+
+
+def main():
+
+    prediction = torch.rand(size=(1, 160000), dtype=torch.float32)
+    ground_truth = torch.rand(size=(1, 160000), dtype=torch.float32)
+
+    prediction_list_r = list(prediction.cpu().numpy())
+    ground_truth_list_r = list(ground_truth.cpu().numpy())
+
+    pesq_score = batch_pesq(prediction_list_r, ground_truth_list_r)
+    print(pesq_score)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/clean_unet/training.py

@@ -0,0 +1,85 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+
+
+def anneal_linear(start, end, proportion):
+    return start + proportion * (end - start)
+
+
+def anneal_cosine(start, end, proportion):
+    cos_val = math.cos(math.pi * proportion) + 1
+    return end + (start - end) / 2 * cos_val
+
+
+class Phase:
+    def __init__(self, start, end, n_iter, cur_iter, anneal_fn):
+        self.start, self.end = start, end
+        self.n_iter = n_iter
+        self.anneal_fn = anneal_fn
+        self.n = cur_iter
+
+    def step(self):
+        self.n += 1
+
+        return self.anneal_fn(self.start, self.end, self.n / self.n_iter)
+
+    def reset(self):
+        self.n = 0
+
+    @property
+    def is_done(self):
+        return self.n >= self.n_iter
+
+
+class LinearWarmupCosineDecay(object):
+    def __init__(
+        self,
+        optimizer,
+        lr_max,
+        n_iter,
+        iteration=0,
+        divider=25,
+        warmup_proportion=0.3,
+        phase=('linear', 'cosine'),
+    ):
+        self.optimizer = optimizer
+
+        phase1 = int(n_iter * warmup_proportion)
+        phase2 = n_iter - phase1
+        lr_min = lr_max / divider
+
+        phase_map = {'linear': anneal_linear, 'cosine': anneal_cosine}
+
+        cur_iter_phase1 = iteration
+        cur_iter_phase2 = max(0, iteration - phase1)
+        self.lr_phase = [
+            Phase(lr_min, lr_max, phase1, cur_iter_phase1, phase_map[phase[0]]),
+            Phase(lr_max, lr_min / 1e4, phase2, cur_iter_phase2, phase_map[phase[1]]),
+        ]
+
+        if iteration < phase1:
+            self.phase = 0
+        else:
+            self.phase = 1
+
+    def step(self):
+        lr = self.lr_phase[self.phase].step()
+
+        for group in self.optimizer.param_groups:
+            group['lr'] = lr
+
+        if self.lr_phase[self.phase].is_done:
+            self.phase += 1
+
+        if self.phase >= len(self.lr_phase):
+            for phase in self.lr_phase:
+                phase.reset()
+
+            self.phase = 0
+
+        return lr
+
+
+if __name__ == '__main__':
+    pass