speaker encoder implementation

2026-09-11 12:00:24 +08:00 · 2019-11-01 12:23:03 +01:00
parent ec579d02a1
commit 6906adc39e
13 changed files with 1179 additions and 0 deletions
@@ -0,0 +1,41 @@
+import os
+import datetime
+import torch
+
+
+def save_checkpoint(model, optimizer, model_loss, out_path,
+                    current_step, epoch):
+    checkpoint_path = 'checkpoint_{}.pth.tar'.format(current_step)
+    checkpoint_path = os.path.join(out_path, checkpoint_path)
+    print(" | | > Checkpoint saving : {}".format(checkpoint_path))
+
+    new_state_dict = model.state_dict()
+    state = {
+        'model': new_state_dict,
+        'optimizer': optimizer.state_dict() if optimizer is not None else None,
+        'step': current_step,
+        'epoch': epoch,
+        'GE2Eloss': model_loss,
+        'date': datetime.date.today().strftime("%B %d, %Y"),
+    }
+    torch.save(state, checkpoint_path)
+
+
+def save_best_model(model, optimizer, model_loss, best_loss, out_path,
+                    current_step):
+    if model_loss < best_loss:
+        new_state_dict = model.state_dict()
+        state = {
+            'model': new_state_dict,
+            'optimizer': optimizer.state_dict(),
+            'step': current_step,
+            'GE2Eloss': model_loss,
+            'date': datetime.date.today().strftime("%B %d, %Y"),
+        }
+        best_loss = model_loss
+        bestmodel_path = 'best_model.pth.tar'
+        bestmodel_path = os.path.join(out_path, bestmodel_path)
+        print("\n > BEST MODEL ({0:.5f}) : {1:}".format(
+            model_loss, bestmodel_path))
+        torch.save(state, bestmodel_path)
+    return best_loss