update

examples/conv_tasnet/step_2_train_model.py

@@ -48,6 +48,8 @@ def get_args():
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
     parser.add_argument("--seed", default=1234, type=int)
 
+    parser.add_argument("--eval_steps", default=5000, type=int)
+
     parser.add_argument("--config_file", default="config.yaml", type=str)
 
 
@@ -237,10 +239,13 @@ def main():
         total_neg_stoi_loss = 0.
         total_mr_stft_loss = 0.
         total_batches = 0.
+        total_steps = 0
         progress_bar = tqdm(
             desc="Training; epoch-{}".format(idx_epoch),
         )
         for batch in train_data_loader:
+            total_steps += 1
+
             clean_audios, noisy_audios = batch
             clean_audios = clean_audios.to(device)
             noisy_audios = noisy_audios.to(device)
@@ -293,128 +298,126 @@ def main():
                 "mr_stft_loss": average_mr_stft_loss,
             })
 
-        # evaluation
-        model.eval()
-        torch.cuda.empty_cache()
-
-        total_pesq_score = 0.
-        total_loss = 0.
-        total_ae_loss = 0.
-        total_neg_si_snr_loss = 0.
-        total_neg_stoi_loss = 0.
-        total_batches = 0.
-
-        progress_bar = tqdm(
-            desc="Evaluation; epoch-{}".format(idx_epoch),
-        )
-        with torch.no_grad():
-            for batch in valid_data_loader:
-                clean_audios, noisy_audios = batch
-                clean_audios = clean_audios.to(device)
-                noisy_audios = noisy_audios.to(device)
-
-                denoise_audios = model.forward(noisy_audios)
-                denoise_audios = torch.squeeze(denoise_audios, dim=1)
-
-                ae_loss = ae_loss_fn.forward(denoise_audios, clean_audios)
-                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)
-
-                # 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
-
-                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")
-
-                total_pesq_score += pesq_score
-                total_loss += loss.item()
-                total_ae_loss += ae_loss.item()
-                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_batches += 1
-
-                average_pesq_score = round(total_pesq_score / total_batches, 4)
-                average_loss = round(total_loss / total_batches, 4)
-                average_ae_loss = round(total_ae_loss / total_batches, 4)
-                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)
-
-                progress_bar.update(1)
-                progress_bar.set_postfix({
-                    "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,
-                    "mr_stft_loss": average_mr_stft_loss,
-
-                })
-
-        # scheduler
-        lr_scheduler.step()
-
-        # save path
-        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
-        epoch_dir.mkdir(parents=True, exist_ok=False)
-
-        # save models
-        model.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 optim
-        torch.save(optimizer.state_dict(), (epoch_dir / "optimizer.pth").as_posix())
-
-        # save metric
-        if best_metric is None:
-            best_idx_epoch = idx_epoch
-            best_metric = average_loss
-        elif average_loss < best_metric:
-            # great is better.
-            best_idx_epoch = idx_epoch
-            best_metric = average_loss
-        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 = 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
+            # evaluation
+            if total_steps % args.eval_steps:
+                model.eval()
+                torch.cuda.empty_cache()
+
+                total_pesq_score = 0.
+                total_loss = 0.
+                total_ae_loss = 0.
+                total_neg_si_snr_loss = 0.
+                total_neg_stoi_loss = 0.
+                total_batches = 0.
+
+                progress_bar = tqdm(
+                    desc="Evaluation; epoch-{}".format(idx_epoch),
+                )
+                with torch.no_grad():
+                    for batch in valid_data_loader:
+                        clean_audios, noisy_audios = batch
+                        clean_audios = clean_audios.to(device)
+                        noisy_audios = noisy_audios.to(device)
+
+                        denoise_audios = model.forward(noisy_audios)
+                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+                        ae_loss = ae_loss_fn.forward(denoise_audios, clean_audios)
+                        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)
+
+                        # 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
+
+                        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")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_ae_loss += ae_loss.item()
+                        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_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_ae_loss = round(total_ae_loss / total_batches, 4)
+                        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)
+
+                        progress_bar.update(1)
+                        progress_bar.set_postfix({
+                            "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,
+                            "mr_stft_loss": average_mr_stft_loss,
+
+                        })
+
+            # save path
+            epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+            epoch_dir.mkdir(parents=True, exist_ok=False)
+
+            # save models
+            model.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 optim
+            torch.save(optimizer.state_dict(), (epoch_dir / "optimizer.pth").as_posix())
+
+            # save metric
+            if best_metric is None:
+                best_idx_epoch = idx_epoch
+                best_metric = average_loss
+            elif average_loss < best_metric:
+                # great is better.
+                best_idx_epoch = idx_epoch
+                best_metric = average_loss
+            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 = 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