update

engine.py

@@ -39,9 +39,7 @@ class YoloTrainer(object):
         self.heavy_eval = False
         # weak augmentatino stage
         self.second_stage = False
-        self.third_stage = False
         self.second_stage_epoch = args.no_aug_epoch
-        self.third_stage_epoch = args.no_aug_epoch // 2
         # path to save model
         self.path_to_save = os.path.join(args.save_folder, args.dataset, args.model)
         os.makedirs(self.path_to_save, exist_ok=True)
@@ -109,20 +107,6 @@ class YoloTrainer(object):
                             'args': self.args}, 
                             checkpoint_path)
 
-            # check third stage
-            if epoch >= (self.args.max_epoch - self.third_stage_epoch - 1) and not self.third_stage:
-                self.check_third_stage()
-                # save model of the last mosaic epoch
-                weight_name = '{}_last_weak_augment_epoch.pth'.format(self.args.model)
-                checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                print('Saving state of the last weak augment epoch-{}.'.format(self.epoch))
-                torch.save({'model': model.state_dict(),
-                            'mAP': round(self.evaluator.map*100, 1),
-                            'optimizer': self.optimizer.state_dict(),
-                            'epoch': self.epoch,
-                            'args': self.args}, 
-                            checkpoint_path)
-
             # train one epoch
             self.epoch = epoch
             self.train_one_epoch(model)
@@ -274,8 +258,7 @@ class YoloTrainer(object):
                 break
 
         # LR Schedule
-        if not self.second_stage:
-            self.lr_scheduler.step()
+        self.lr_scheduler.step()
 
         # Gather the stats from all processes
         metric_logger.synchronize_between_processes()
@@ -312,7 +295,6 @@ class YoloTrainer(object):
 
         # Choose a new image size
         new_img_size = random.randrange(min_img_size, max_img_size + max_stride, max_stride)
-        new_img_size = new_img_size // max_stride * max_stride  # size
 
         if new_img_size / old_img_size != 1:
             # interpolate
@@ -373,25 +355,6 @@ class YoloTrainer(object):
             args=self.args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=True)
         self.train_loader.dataset.transform = self.train_transform
         
-    def check_third_stage(self):
-        # set third stage
-        print('============== Third stage of Training ==============')
-        self.third_stage = True
-
-        # close random affine
-        if 'translate' in self.trans_cfg.keys() and self.trans_cfg['translate'] > 0.0:
-            print(' - Close < translate of affine > ...')
-            self.trans_cfg['translate'] = 0.0
-        if 'scale' in self.trans_cfg.keys():
-            print(' - Close < scale of affine >...')
-            self.trans_cfg['scale'] = [1.0, 1.0]
-
-        # build a new transform for second stage
-        print(' - Rebuild transforms ...')
-        self.train_transform, self.trans_cfg = build_transform(
-            args=self.args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=True)
-        self.train_loader.dataset.transform = self.train_transform
-
 ## Customed Trainer for YOLOX series
 class YoloxTrainer(object):
     def __init__(self, args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size):