From 7ea2007326056a371556ad2deaaf1146eea4bce0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 31 Mar 2024 06:14:22 +0200 Subject: [PATCH] Favor EMA over `model` in train checkpoints (#9433) Signed-off-by: Glenn Jocher --- ultralytics/engine/trainer.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index 2dcbf9fe..7a963a38 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -42,7 +42,6 @@ from ultralytics.utils.files import get_latest_run from ultralytics.utils.torch_utils import ( EarlyStopping, ModelEMA, - de_parallel, init_seeds, one_cycle, select_device, @@ -486,7 +485,7 @@ class BaseTrainer: { "epoch": self.epoch, "best_fitness": self.best_fitness, - "model": deepcopy(de_parallel(self.model)).half(), + "model": None, # resume and final checkpoints derive from EMA "ema": deepcopy(self.ema.ema).half(), "updates": self.ema.updates, "optimizer": self.optimizer.state_dict(), @@ -527,7 +526,7 @@ class BaseTrainer: ckpt = None if str(model).endswith(".pt"): weights, ckpt = attempt_load_one_weight(model) - cfg = ckpt["model"].yaml + cfg = weights.yaml else: cfg = model self.model = self.get_model(cfg=cfg, weights=weights, verbose=RANK == -1) # calls Model(cfg, weights)