update

examples/mpnet_aishell/step_2_train_model.py

@@ -167,18 +167,18 @@ def main():
     discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
 
     # resume training
-    epoch_max = -1
+    last_epoch = -1
     for epoch_i in serialization_dir.glob("epoch-*"):
         epoch_i = Path(epoch_i)
         epoch_idx = epoch_i.stem.split("-")[1]
         epoch_idx = int(epoch_idx)
-        if epoch_idx > epoch_max:
-            epoch_max = epoch_idx
+        if epoch_idx > last_epoch:
+            last_epoch = epoch_idx
 
-    if epoch_max != -1:
-        logger.info(f"resume from epoch-{epoch_max}.")
-        generator_pt = serialization_dir / f"epoch-{epoch_max}/generator.pt"
-        discriminator_pt = serialization_dir / f"epoch-{epoch_max}/discriminator.pt"
+    if last_epoch != -1:
+        logger.info(f"resume from epoch-{last_epoch}.")
+        generator_pt = serialization_dir / f"epoch-{last_epoch}/generator.pt"
+        discriminator_pt = serialization_dir / f"epoch-{last_epoch}/discriminator.pt"
 
         logger.info(f"load state dict for generator.")
         with open(generator_pt.as_posix(), "rb") as f:
@@ -199,8 +199,8 @@ def main():
     optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
     optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
 
-    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=-1)
-    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=-1)
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
 
     # training loop
 
@@ -219,7 +219,7 @@ def main():
     patience_count = 0
 
     logger.info("training")
-    for idx_epoch in range(args.max_epochs):
+    for idx_epoch in range(max(0, last_epoch), args.max_epochs):
         # train
         generator.train()
         discriminator.train()