yjh0410 1 سال پیش
والد
کامیت
140812e3f7
2فایلهای تغییر یافته به همراه8 افزوده شده و 0 حذف شده
  1. 2 0
      yolo/engine.py
  2. 6 0
      yolo/utils/solver/optimizer.py

+ 2 - 0
yolo/engine.py

@@ -72,6 +72,8 @@ class YoloTrainer(object):
         self.lr_scheduler_warmup = LinearWarmUpLrScheduler(warmup_iters, cfg.base_lr, cfg.warmup_bias_lr)
         self.lr_scheduler = build_lr_scheduler(cfg, self.optimizer, args.resume)
 
+        self.best_map = cfg.best_map
+
     def train(self, model):
         for epoch in range(self.start_epoch, self.cfg.max_epoch):
             if self.args.distributed:

+ 6 - 0
yolo/utils/solver/optimizer.py

@@ -44,6 +44,12 @@ def build_yolo_optimizer(cfg, model, resume=None):
             optimizer.load_state_dict(checkpoint_state_dict)
             start_epoch = checkpoint.pop("epoch") + 1
             del checkpoint, checkpoint_state_dict
+            if "mAP" in checkpoint:
+                print('--Load best metric from the checkpoint: ', resume)
+                best_map = checkpoint["mAP"]
+                cfg.best_map = best_map
+            else:
+                cfg.best_map = -1.
         except:
             print("No optimzier in the given checkpoint.")