update

examples/conv_tasnet/step_2_train_model.py

@@ -31,7 +31,6 @@ import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from torch.utils.data.dataloader import DataLoader
-from torch_pesq import PesqLoss
 from tqdm import tqdm
 
 from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
@@ -39,7 +38,7 @@ from toolbox.torchaudio.models.conv_tasnet.configuration_conv_tasnet import Conv
 from toolbox.torchaudio.models.conv_tasnet.modeling_conv_tasnet import ConvTasNet, ConvTasNetPretrainedModel
 from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
 from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
-from toolbox.torchaudio.losses.perceptual import NegSTOILoss
+from toolbox.torchaudio.losses.perceptual import NegSTOILoss, PesqLoss
 from toolbox.torchaudio.metrics.pesq import run_pesq_score
 
 
@@ -283,7 +282,6 @@ def main():
             neg_stoi_loss = neg_stoi_loss_fn.forward(denoise_audios, clean_audios)
             mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
             pesq_loss = pesq_loss_fn.forward(clean_audios, denoise_audios)
-            print(f"pesq_loss: {pesq_loss}")
 
             # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
             # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss

toolbox/torchaudio/losses/perceptual.py

@@ -6,6 +6,7 @@ https://zhuanlan.zhihu.com/p/627039860
 import torch
 import torch.nn as nn
 from torch_stoi import NegSTOILoss as TorchNegSTOILoss
+from torch_pesq import PesqLoss as TorchPesqLoss
 
 
 class PMSQELoss(object):
@@ -55,6 +56,47 @@ class NegSTOILoss(nn.Module):
         return loss
 
 
+class PesqLoss(nn.Module):
+    def __init__(self,
+                 factor: float,
+                 sample_rate: int = 48000,
+                 nbarks: int = 49,
+                 win_length: int = 512,
+                 n_fft: int = 512,
+                 hop_length: int = 256,
+                 reduction: str = "mean",
+                 ):
+        super(PesqLoss, self).__init__()
+        self.factor = factor
+        self.sample_rate = sample_rate
+        self.nbarks = nbarks
+        self.win_length = win_length
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.reduction = reduction
+
+        self.loss_fn = TorchPesqLoss(
+            factor=factor,
+            sample_rate=sample_rate,
+            nbarks=nbarks,
+            win_length=win_length,
+            n_fft=n_fft,
+            hop_length=hop_length,
+        )
+
+    def forward(self, denoise: torch.Tensor, clean: torch.Tensor):
+
+        batch_loss = self.loss_fn.forward(clean, denoise)
+
+        if self.reduction == "mean":
+            loss = torch.mean(batch_loss)
+        elif self.reduction == "sum":
+            loss = torch.sum(batch_loss)
+        else:
+            raise AssertionError
+        return loss
+
+
 def main():
     sample_rate = 16000