update

examples/conv_tasnet/step_2_train_model.py

@@ -2,6 +2,14 @@
 # -*- coding: utf-8 -*-
 """
 https://github.com/kaituoxu/Conv-TasNet/tree/master/src
+
+一般场景：
+
+目标 SI-SNR ≥ 10 dB，适用于电话通信、基础语音助手等。
+
+高要求场景（如医疗助听、语音识别）：
+需 SI-SNR ≥ 14 dB，并配合 PESQ ≥ 3.0 和 STOI ≥ 0.851812。
+
 """
 import argparse
 import json
@@ -23,6 +31,7 @@ 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
@@ -129,12 +138,12 @@ def main():
     # datasets
     train_dataset = DenoiseJsonlDataset(
         jsonl_file=args.train_dataset,
-        expected_sample_rate=8000,
+        expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
     )
     valid_dataset = DenoiseJsonlDataset(
         jsonl_file=args.valid_dataset,
-        expected_sample_rate=8000,
+        expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
     )
     train_data_loader = DataLoader(
@@ -213,7 +222,7 @@ def main():
 
     ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
     neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
-    neg_stoi_loss_fn = NegSTOILoss(sample_rate=8000, reduction="mean").to(device)
+    neg_stoi_loss_fn = NegSTOILoss(sample_rate=config.sample_rate, reduction="mean").to(device)
     mr_stft_loss_fn = MultiResolutionSTFTLoss(
         fft_size_list=[256, 512, 1024],
         win_size_list=[120, 240, 480],
@@ -222,6 +231,7 @@ def main():
         factor_mag=1.0,
         reduction="mean"
     ).to(device)
+    pesq_loss_fn = PesqLoss(0.5, sample_rate=config.sample_rate).to(device)
 
     # training loop
 
@@ -249,6 +259,7 @@ def main():
         total_neg_si_snr_loss = 0.
         total_neg_stoi_loss = 0.
         total_mr_stft_loss = 0.
+        total_pesq_loss = 0.
         total_batches = 0.
 
         step_idx = 0 if last_step_idx == -1 else last_step_idx
@@ -271,16 +282,18 @@ def main():
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
             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)
 
             # 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
             # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
             # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
-            loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+            # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+            loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=8000, mode="nb")
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
 
             optimizer.zero_grad()
             loss.backward()
@@ -293,6 +306,7 @@ def main():
             total_neg_si_snr_loss += neg_si_snr_loss.item()
             total_neg_stoi_loss += neg_stoi_loss.item()
             total_mr_stft_loss += mr_stft_loss.item()
+            total_pesq_loss += pesq_loss.item()
             total_batches += 1
 
             average_pesq_score = round(total_pesq_score / total_batches, 4)
@@ -301,6 +315,7 @@ def main():
             average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
             average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
             average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_pesq_loss = round(total_pesq_loss / total_batches, 4)
 
             progress_bar_train.update(1)
             progress_bar_train.set_postfix({
@@ -311,6 +326,7 @@ def main():
                 "neg_si_snr_loss": average_neg_si_snr_loss,
                 "neg_stoi_loss": average_neg_stoi_loss,
                 "mr_stft_loss": average_mr_stft_loss,
+                "pesq_loss": average_pesq_loss,
             })
 
             # evaluation
@@ -325,6 +341,7 @@ def main():
                     total_neg_si_snr_loss = 0.
                     total_neg_stoi_loss = 0.
                     total_mr_stft_loss = 0.
+                    total_pesq_loss = 0.
                     total_batches = 0.
 
                     progress_bar_train.close()
@@ -343,16 +360,18 @@ def main():
                         neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
                         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)
 
                         # 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
                         # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
                         # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
-                        loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+                        # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+                        loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
 
                         denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
                         clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=8000, mode="nb")
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
 
                         total_pesq_score += pesq_score
                         total_loss += loss.item()
@@ -360,6 +379,7 @@ def main():
                         total_neg_si_snr_loss += neg_si_snr_loss.item()
                         total_neg_stoi_loss += neg_stoi_loss.item()
                         total_mr_stft_loss += mr_stft_loss.item()
+                        total_pesq_loss += pesq_loss.item()
                         total_batches += 1
 
                         average_pesq_score = round(total_pesq_score / total_batches, 4)
@@ -368,6 +388,7 @@ def main():
                         average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
                         average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
                         average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_pesq_loss = round(total_pesq_loss / total_batches, 4)
 
                         progress_bar_eval.update(1)
                         progress_bar_eval.set_postfix({
@@ -378,6 +399,7 @@ def main():
                             "neg_si_snr_loss": average_neg_si_snr_loss,
                             "neg_stoi_loss": average_neg_stoi_loss,
                             "mr_stft_loss": average_mr_stft_loss,
+                            "pesq_loss": average_pesq_loss,
                         })
 
                     total_pesq_score = 0.
@@ -386,6 +408,7 @@ def main():
                     total_neg_si_snr_loss = 0.
                     total_neg_stoi_loss = 0.
                     total_mr_stft_loss = 0.
+                    total_pesq_loss = 0.
                     total_batches = 0.
 
                     progress_bar_eval.close()