Merge pull request #201 from supervisely-ecosystem/best_epoch

added best epoch number to best checkpoint file
supervisely-ecosystem · Jun 16, 2023 · 36bf76b · 36bf76b
2 parents e5af1e0 + fb902e2
commit 36bf76b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -67,7 +67,6 @@
 
 
 def train(hyp, opt, device, tb_writer=None):
-
     train_batches_uploaded = False
 
     logger.info("hyperparameters", extra=hyp)
@@ -626,6 +625,9 @@ def train(hyp, opt, device, tb_writer=None):
                 # Save last, best and delete
                 torch.save(ckpt, last)
                 if best_fitness == fi:
+                    if epoch != 0:
+                        os.remove(best)
+                    best = wdir / f"best_{epoch}.pt"
                     torch.save(ckpt, best)
                 if wandb_logger.wandb:
                     if (