save mosaic model

exp-a
Laughing-q 1 year ago
parent 148e28c60f
commit 6256e711ed
  1. 4
      ultralytics/engine/trainer.py

@ -112,6 +112,7 @@ class BaseTrainer:
self.args.save_dir = str(self.save_dir)
yaml_save(self.save_dir / "args.yaml", vars(self.args)) # save run args
self.last, self.best = self.wdir / "last.pt", self.wdir / "best.pt" # checkpoint paths
self.last_mosaic = self.wdir / "last_mosaic.pt"
self.save_period = self.args.save_period
self.batch_size = self.args.batch
@ -490,6 +491,9 @@ class BaseTrainer:
"version": __version__,
}
if self.args.close_mosaic and self.epoch == (self.epochs - self.args.close_mosaic - 1):
torch.save(ckpt, self.last_mosaic)
# Save last and best
torch.save(ckpt, self.last)
if self.best_fitness == self.fitness:

Loading…
Cancel
Save