[train] fix epoch when loading from previous checkpoint (#2252)

wenet-e2e · Dec 19, 2023 · 54c07fa · 54c07fa
1 parent 56c0e8e
commit 54c07fa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/wenet/bin/train.py b/wenet/bin/train.py
@@ -122,7 +122,9 @@ def main():
         scaler = torch.cuda.amp.GradScaler()
 
     # Start training loop
-    start_epoch = configs["init_infos"].get('epoch', -1) + 1
+    tag = configs["init_infos"].get("tag", "init")
+    start_epoch = configs["init_infos"].get('epoch', -1) + int("epoch_" in tag)
+    configs.pop("init_infos", None)
     final_epoch = None
     for epoch in range(start_epoch, configs.get('max_epoch', 100)):
         train_dataset.set_epoch(epoch)

diff --git a/wenet/utils/executor.py b/wenet/utils/executor.py
@@ -83,6 +83,7 @@ def train(self, model, optimizer, scheduler, train_data_loader,
                         and (batch_idx + 1) % info_dict["accum_grad"] == 0:
                     total_loss, num_seen_utts = self.cv(
                         model, cv_data_loader, configs)
+                    model.train()
                     info_dict.update({
                         "tag":
                         "step_{}".format(self.step),