Merge pull request meituan#794 from yuyang-cloud/main

fix resume bug of mAP drop.
ZHEQIUSHUI · Apr 25, 2023 · d9c0491 · d9c0491
2 parents 81a9580 + f6565fb
commit d9c0491
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/yolov6/core/engine.py b/yolov6/core/engine.py
@@ -79,6 +79,7 @@ def __init__(self, args, cfg, device):
             model.load_state_dict(resume_state_dict, strict=True)  # load
             self.start_epoch = self.ckpt['epoch'] + 1
             self.optimizer.load_state_dict(self.ckpt['optimizer'])
+            self.scheduler.load_state_dict(self.ckpt['scheduler'])
             if self.main_process:
                 self.ema.ema.load_state_dict(self.ckpt['ema'].float().state_dict())
                 self.ema.updates = self.ckpt['updates']
@@ -179,6 +180,7 @@ def eval_and_save(self):
                     'ema': deepcopy(self.ema.ema).half(),
                     'updates': self.ema.updates,
                     'optimizer': self.optimizer.state_dict(),
+                    'scheduler': self.scheduler.state_dict(),
                     'epoch': self.epoch,
                     'results': self.evaluate_results,
                     }
@@ -301,6 +303,9 @@ def train_before_loop(self):
     def prepare_for_steps(self):
         if self.epoch > self.start_epoch:
             self.scheduler.step()
+        elif  hasattr(self, "ckpt") and self.epoch == self.start_epoch: # resume first epoch, load lr
+            for k, param in enumerate(self.optimizer.param_groups):
+                param['lr'] = self.scheduler.get_lr()[k]
         #stop strong aug like mosaic and mixup from last n epoch by recreate dataloader
         if self.epoch == self.max_epoch - self.args.stop_aug_last_n_epoch:
             self.cfg.data_aug.mosaic = 0.0