From 6256e711ed7475f858171860db4b601c961ac8bc Mon Sep 17 00:00:00 2001
From: Laughing-q <1185102784@qq.com>
Date: Thu, 29 Feb 2024 19:14:03 +0800
Subject: [PATCH] save mosaic model

---
 ultralytics/engine/trainer.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index 33821171e4..33c9b0e002 100644
--- a/ultralytics/engine/trainer.py
+++ b/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: