update

examples/conv_tasnet/step_2_train_model.py

@@ -359,64 +359,64 @@ def main():
 
                         })
 
-            # save path
-            save_dir = serialization_dir / "steps-{}".format(total_steps)
-            save_dir.mkdir(parents=True, exist_ok=False)
-
-            # save models
-            model.save_pretrained(save_dir.as_posix())
-
-            model_list.append(save_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 optim
-            torch.save(optimizer.state_dict(), (save_dir / "optimizer.pth").as_posix())
-
-            # save metric
-            if best_metric is None:
-                best_idx_epoch = idx_epoch
-                best_metric = average_pesq_score
-            elif average_pesq_score > best_metric:
-                # great is better.
-                best_idx_epoch = idx_epoch
-                best_metric = average_pesq_score
-            else:
-                pass
-
-            metrics = {
-                "idx_epoch": idx_epoch,
-                "best_idx_epoch": best_idx_epoch,
-                "pesq_score": average_pesq_score,
-                "loss": average_loss,
-                "ae_loss": average_ae_loss,
-                "neg_si_snr_loss": average_neg_si_snr_loss,
-                "neg_stoi_loss": average_neg_stoi_loss,
-            }
-            metrics_filename = save_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(save_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
+                # save path
+                save_dir = serialization_dir / "steps-{}".format(total_steps)
+                save_dir.mkdir(parents=True, exist_ok=False)
+
+                # save models
+                model.save_pretrained(save_dir.as_posix())
+
+                model_list.append(save_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 optim
+                torch.save(optimizer.state_dict(), (save_dir / "optimizer.pth").as_posix())
+
+                # save metric
+                if best_metric is None:
+                    best_idx_epoch = idx_epoch
+                    best_metric = average_pesq_score
+                elif average_pesq_score > best_metric:
+                    # great is better.
+                    best_idx_epoch = idx_epoch
+                    best_metric = average_pesq_score
+                else:
+                    pass
+
+                metrics = {
+                    "idx_epoch": idx_epoch,
+                    "best_idx_epoch": best_idx_epoch,
+                    "pesq_score": average_pesq_score,
+                    "loss": average_loss,
+                    "ae_loss": average_ae_loss,
+                    "neg_si_snr_loss": average_neg_si_snr_loss,
+                    "neg_stoi_loss": average_neg_stoi_loss,
+                }
+                metrics_filename = save_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(save_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
 

toolbox/torchaudio/losses/spectral.py

@@ -39,8 +39,8 @@ class LSDLoss(nn.Module):
 
     def forward(self, denoise_power: torch.Tensor, clean_power: torch.Tensor):
         """
-        :param denoise_power: The estimated signal (batch_size, signal_length)
-        :param clean_power: The target signal (batch_size, signal_length)
+        :param denoise_power: power spectrum of the estimated signal power spectrum  (batch_size, ...)
+        :param clean_power: power spectrum of the target signal (batch_size, ...)
         :return:
         """
         denoise_power = denoise_power + self.eps