yjh0410 2 éve
szülő
commit
d48ba0c4cc
2 módosított fájl, 5 hozzáadás és 5 törlés
  1. 4 4
      engine.py
  2. 1 1
      utils/solver/optimizer.py

+ 4 - 4
engine.py

@@ -79,7 +79,7 @@ class Yolov8Trainer(object):
 
         # ---------------------------- Build LR Scheduler ----------------------------
         self.lr_scheduler, self.lf = build_lr_scheduler(self.lr_schedule_dict, self.optimizer, self.args.max_epoch)
-        self.lr_scheduler.last_epoch = self.start_epoch - 1  # do not move
+        self.lr_scheduler.last_epoch = self.start_epoch - 2  # do not move
         if self.args.resume and self.args.resume != 'None':
             self.lr_scheduler.step()
 
@@ -449,7 +449,7 @@ class YoloxTrainer(object):
 
         # ---------------------------- Build LR Scheduler ----------------------------
         self.lr_scheduler, self.lf = build_lr_scheduler(self.lr_schedule_dict, self.optimizer, self.args.max_epoch - self.no_aug_epoch)
-        self.lr_scheduler.last_epoch = self.start_epoch - 1  # do not move
+        self.lr_scheduler.last_epoch = self.start_epoch - 2  # do not move
         if self.args.resume and self.args.resume != 'None':
             self.lr_scheduler.step()
 
@@ -814,7 +814,7 @@ class RTCTrainer(object):
 
         # ---------------------------- Build LR Scheduler ----------------------------
         self.lr_scheduler, self.lf = build_lr_scheduler(self.lr_schedule_dict, self.optimizer, args.max_epoch - args.no_aug_epoch)
-        self.lr_scheduler.last_epoch = self.start_epoch - 1  # do not move
+        self.lr_scheduler.last_epoch = self.start_epoch - 2  # do not move
         if self.args.resume and self.args.resume != 'None':
             self.lr_scheduler.step()
 
@@ -1185,7 +1185,7 @@ class RTRTrainer(object):
 
         # ---------------------------- Build LR Scheduler ----------------------------
         self.lr_scheduler, self.lf = build_lr_scheduler(self.lr_schedule_dict, self.optimizer, args.max_epoch - args.no_aug_epoch)
-        self.lr_scheduler.last_epoch = self.start_epoch - 1  # do not move
+        self.lr_scheduler.last_epoch = self.start_epoch - 2  # do not move
         if self.args.resume and self.args.resume != 'None':
             self.lr_scheduler.step()
 

+ 1 - 1
utils/solver/optimizer.py

@@ -34,7 +34,7 @@ def build_yolo_optimizer(cfg, model, resume=None):
     start_epoch = 0
     if resume and resume != "None":
         print('keep training: ', resume)
-        checkpoint = torch.load(resume)
+        checkpoint = torch.load(resume, map_location='cpu')
         # checkpoint state dict
         checkpoint_state_dict = checkpoint.pop("optimizer")
         optimizer.load_state_dict(checkpoint_state_dict)