use 6 decoder layers for R18

config/model_config/rtdetr_config.py

@@ -30,7 +30,7 @@ rtdetr_cfg = {
         # Transformer Decoder
         'transformer': 'rtdetr_transformer',
         'de_num_heads': 8,
-        'de_num_layers': 3,
+        'de_num_layers': 6,
         'de_ffn_dim': 1024,
         'de_dropout': 0.0,
         'de_act': 'relu',
@@ -52,7 +52,7 @@ rtdetr_cfg = {
                        'giou': 2,},
         # ---------------- Train config ----------------
         ## input
-        'multi_scale': [0.5, 1.25],   # 320 -> 800
+        'multi_scale': [0.5, 1.5],   # 320 -> 960
         'trans_type': 'rtdetr_s',
         # ---------------- Train config ----------------
         'trainer_type': 'rtdetr',

engine.py

@@ -1598,693 +1598,6 @@ class RTPDetrTrainer(RTDetrTrainer):
             self.lr_scheduler.step()
         
 
-# ----------------------- Det + Seg trainers -----------------------
-## RTCDet Trainer for Det + Seg
-class RTCTrainerDS(object):
-    def __init__(self, args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size):
-        # ------------------- basic parameters -------------------
-        self.args = args
-        self.epoch = 0
-        self.best_map = -1.
-        self.device = device
-        self.criterion = criterion
-        self.world_size = world_size
-        self.grad_accumulate = args.grad_accumulate
-        self.clip_grad = 35
-        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)
-
-        # ---------------------------- Hyperparameters refer to RTMDet ----------------------------
-        self.optimizer_dict = {'optimizer': 'adamw', 'momentum': None, 'weight_decay': 5e-2, 'lr0': 0.001}
-        self.ema_dict = {'ema_decay': 0.9998, 'ema_tau': 2000}
-        self.lr_schedule_dict = {'scheduler': 'linear', 'lrf': 0.01}
-        self.warmup_dict = {'warmup_momentum': 0.8, 'warmup_bias_lr': 0.1}        
-
-        # ---------------------------- Build Dataset & Model & Trans. Config ----------------------------
-        self.data_cfg = data_cfg
-        self.model_cfg = model_cfg
-        self.trans_cfg = trans_cfg
-
-        # ---------------------------- Build Transform ----------------------------
-        self.train_transform, self.trans_cfg = build_transform(
-            args=args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=True)
-        self.val_transform, _ = build_transform(
-            args=args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=False)
-
-        # ---------------------------- Build Dataset & Dataloader ----------------------------
-        self.dataset, self.dataset_info = build_dataset(args, self.data_cfg, self.trans_cfg, self.train_transform, is_train=True)
-        self.train_loader = build_dataloader(args, self.dataset, self.args.batch_size // self.world_size, CollateFunc())
-
-        # ---------------------------- Build Evaluator ----------------------------
-        self.evaluator = build_evluator(args, self.data_cfg, self.val_transform, self.device)
-
-        # ---------------------------- Build Grad. Scaler ----------------------------
-        self.scaler = torch.cuda.amp.GradScaler(enabled=self.args.fp16)
-
-        # ---------------------------- Build Optimizer ----------------------------
-        self.optimizer_dict['lr0'] *= args.batch_size * self.grad_accumulate / 64
-        self.optimizer, self.start_epoch = build_yolo_optimizer(self.optimizer_dict, model, args.resume)
-
-        # ---------------------------- Build LR Scheduler ----------------------------
-        self.lr_scheduler, self.lf = build_lambda_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
-        if self.args.resume and self.args.resume != 'None':
-            self.lr_scheduler.step()
-
-        # ---------------------------- Build Model-EMA ----------------------------
-        if self.args.ema and distributed_utils.get_rank() in [-1, 0]:
-            print('Build ModelEMA ...')
-            self.model_ema = ModelEMA(self.ema_dict, model, self.start_epoch * len(self.train_loader))
-        else:
-            self.model_ema = None
-
-    def train(self, model):
-        for epoch in range(self.start_epoch, self.args.max_epoch):
-            if self.args.distributed:
-                self.train_loader.batch_sampler.sampler.set_epoch(epoch)
-
-            # check second stage
-            if epoch >= (self.args.max_epoch - self.second_stage_epoch - 1) and not self.second_stage:
-                self.check_second_stage()
-                # save model of the last mosaic epoch
-                weight_name = '{}_last_mosaic_epoch.pth'.format(self.args.model)
-                checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                print('Saving state of the last Mosaic 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)
-
-            # 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)
-
-            # eval one epoch
-            if self.heavy_eval:
-                model_eval = model.module if self.args.distributed else model
-                self.eval(model_eval)
-            else:
-                model_eval = model.module if self.args.distributed else model
-                if (epoch % self.args.eval_epoch) == 0 or (epoch == self.args.max_epoch - 1):
-                    self.eval(model_eval)
-
-            if self.args.debug:
-                print("For debug mode, we only train 1 epoch")
-                break
-
-    def eval(self, model):
-        # chech model
-        model_eval = model if self.model_ema is None else self.model_ema.ema
-
-        if distributed_utils.is_main_process():
-            # check evaluator
-            if self.evaluator is None:
-                print('No evaluator ... save model and go on training.')
-                print('Saving state, epoch: {}'.format(self.epoch))
-                weight_name = '{}_no_eval.pth'.format(self.args.model)
-                checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                torch.save({'model': model_eval.state_dict(),
-                            'mAP': -1.,
-                            'optimizer': self.optimizer.state_dict(),
-                            'epoch': self.epoch,
-                            'args': self.args}, 
-                            checkpoint_path)               
-            else:
-                print('eval ...')
-                # set eval mode
-                model_eval.trainable = False
-                model_eval.eval()
-
-                # evaluate
-                with torch.no_grad():
-                    self.evaluator.evaluate(model_eval)
-
-                # save model
-                cur_map = self.evaluator.map
-                if cur_map > self.best_map:
-                    # update best-map
-                    self.best_map = cur_map
-                    # save model
-                    print('Saving state, epoch:', self.epoch)
-                    weight_name = '{}_best.pth'.format(self.args.model)
-                    checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                    torch.save({'model': model_eval.state_dict(),
-                                'mAP': round(self.best_map*100, 1),
-                                'optimizer': self.optimizer.state_dict(),
-                                'epoch': self.epoch,
-                                'args': self.args}, 
-                                checkpoint_path)                      
-
-                # set train mode.
-                model_eval.trainable = True
-                model_eval.train()
-
-        if self.args.distributed:
-            # wait for all processes to synchronize
-            dist.barrier()
-
-    def train_one_epoch(self, model):
-        metric_logger = MetricLogger(delimiter="  ")
-        metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))
-        metric_logger.add_meter('size', SmoothedValue(window_size=1, fmt='{value:d}'))
-        header = 'Epoch: [{} / {}]'.format(self.epoch, self.args.max_epoch)
-        epoch_size = len(self.train_loader)
-        print_freq = 10
-
-        # basic parameters
-        epoch_size = len(self.train_loader)
-        img_size = self.args.img_size
-        nw = epoch_size * self.args.wp_epoch
-
-        # Train one epoch
-        for iter_i, (images, targets) in enumerate(metric_logger.log_every(self.train_loader, print_freq, header)):
-            ni = iter_i + self.epoch * epoch_size
-            # Warmup
-            if ni <= nw:
-                xi = [0, nw]  # x interp
-                for j, x in enumerate(self.optimizer.param_groups):
-                    # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
-                    x['lr'] = np.interp(
-                        ni, xi, [self.warmup_dict['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * self.lf(self.epoch)])
-                    if 'momentum' in x:
-                        x['momentum'] = np.interp(ni, xi, [self.warmup_dict['warmup_momentum'], self.optimizer_dict['momentum']])
-                                
-            # To device
-            images = images.to(self.device, non_blocking=True).float()
-
-            # Multi scale
-            if self.args.multi_scale:
-                images, targets, img_size = self.rescale_image_targets(
-                    images, targets, self.model_cfg['stride'], self.args.min_box_size, self.model_cfg['multi_scale'])
-            else:
-                targets = self.refine_targets(targets, self.args.min_box_size)
-                
-            # Visualize train targets
-            if self.args.vis_tgt:
-                vis_data(images*255, targets, self.data_cfg['num_classes'])
-
-            # Inference
-            with torch.cuda.amp.autocast(enabled=self.args.fp16):
-                outputs = model(images)
-                # Compute loss
-                loss_dict = self.criterion(outputs=outputs, targets=targets, epoch=self.epoch, task='det_seg')
-                det_loss_dict = loss_dict['det_loss_dict']
-                seg_loss_dict = loss_dict['seg_loss_dict']
-
-                # TODO: finish the backward + optimize
-
-            # # Update log
-            # metric_logger.update(**loss_dict_reduced)
-            # metric_logger.update(lr=self.optimizer.param_groups[2]["lr"])
-            # metric_logger.update(grad_norm=grad_norm)
-            # metric_logger.update(size=img_size)
-
-            if self.args.debug:
-                print("For debug mode, we only train 1 iteration")
-                break
-
-        # LR Schedule
-        if not self.second_stage:
-            self.lr_scheduler.step()
-
-        # Gather the stats from all processes
-        metric_logger.synchronize_between_processes()
-        print("Averaged stats:", metric_logger)
-
-    def refine_targets(self, targets, min_box_size):
-        # rescale targets
-        for tgt in targets:
-            boxes = tgt["boxes"].clone()
-            labels = tgt["labels"].clone()
-            # refine tgt
-            tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
-            min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
-            keep = (min_tgt_size >= min_box_size)
-
-            tgt["boxes"] = boxes[keep]
-            tgt["labels"] = labels[keep]
-        
-        return targets
-
-    def rescale_image_targets(self, images, targets, stride, min_box_size, multi_scale_range=[0.5, 1.5]):
-        """
-            Deployed for Multi scale trick.
-        """
-        if isinstance(stride, int):
-            max_stride = stride
-        elif isinstance(stride, list):
-            max_stride = max(stride)
-
-        # During training phase, the shape of input image is square.
-        old_img_size = images.shape[-1]
-        new_img_size = random.randrange(old_img_size * multi_scale_range[0], old_img_size * multi_scale_range[1] + max_stride)
-        new_img_size = new_img_size // max_stride * max_stride  # size
-        if new_img_size / old_img_size != 1:
-            # interpolate
-            images = torch.nn.functional.interpolate(
-                                input=images, 
-                                size=new_img_size, 
-                                mode='bilinear', 
-                                align_corners=False)
-        # rescale targets
-        for tgt in targets:
-            boxes = tgt["boxes"].clone()
-            labels = tgt["labels"].clone()
-            boxes = torch.clamp(boxes, 0, old_img_size)
-            # rescale box
-            boxes[:, [0, 2]] = boxes[:, [0, 2]] / old_img_size * new_img_size
-            boxes[:, [1, 3]] = boxes[:, [1, 3]] / old_img_size * new_img_size
-            # refine tgt
-            tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
-            min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
-            keep = (min_tgt_size >= min_box_size)
-
-            tgt["boxes"] = boxes[keep]
-            tgt["labels"] = labels[keep]
-
-        return images, targets, new_img_size
-
-    def check_second_stage(self):
-        # set second stage
-        print('============== Second stage of Training ==============')
-        self.second_stage = True
-
-        # close mosaic augmentation
-        if self.train_loader.dataset.mosaic_prob > 0.:
-            print(' - Close < Mosaic Augmentation > ...')
-            self.train_loader.dataset.mosaic_prob = 0.
-            self.heavy_eval = True
-
-        # close mixup augmentation
-        if self.train_loader.dataset.mixup_prob > 0.:
-            print(' - Close < Mixup Augmentation > ...')
-            self.train_loader.dataset.mixup_prob = 0.
-            self.heavy_eval = True
-
-        # close rotation augmentation
-        if 'degrees' in self.trans_cfg.keys() and self.trans_cfg['degrees'] > 0.0:
-            print(' - Close < degress of rotation > ...')
-            self.trans_cfg['degrees'] = 0.0
-        if 'shear' in self.trans_cfg.keys() and self.trans_cfg['shear'] > 0.0:
-            print(' - Close < shear of rotation >...')
-            self.trans_cfg['shear'] = 0.0
-        if 'perspective' in self.trans_cfg.keys() and self.trans_cfg['perspective'] > 0.0:
-            print(' - Close < perspective of rotation > ...')
-            self.trans_cfg['perspective'] = 0.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
-        
-    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
-   
-
-# ----------------------- Det + Seg + Pos trainers -----------------------
-## RTCDet Trainer for Det + Seg + HumanPose
-class RTCTrainerDSP(object):
-    def __init__(self, args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size):
-        # ------------------- basic parameters -------------------
-        self.args = args
-        self.epoch = 0
-        self.best_map = -1.
-        self.device = device
-        self.criterion = criterion
-        self.world_size = world_size
-        self.grad_accumulate = args.grad_accumulate
-        self.clip_grad = 35
-        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)
-
-        # ---------------------------- Hyperparameters refer to RTMDet ----------------------------
-        self.optimizer_dict = {'optimizer': 'adamw', 'momentum': None, 'weight_decay': 5e-2, 'lr0': 0.001}
-        self.ema_dict = {'ema_decay': 0.9998, 'ema_tau': 2000}
-        self.lr_schedule_dict = {'scheduler': 'linear', 'lrf': 0.01}
-        self.warmup_dict = {'warmup_momentum': 0.8, 'warmup_bias_lr': 0.1}        
-
-        # ---------------------------- Build Dataset & Model & Trans. Config ----------------------------
-        self.data_cfg = data_cfg
-        self.model_cfg = model_cfg
-        self.trans_cfg = trans_cfg
-
-        # ---------------------------- Build Transform ----------------------------
-        self.train_transform, self.trans_cfg = build_transform(
-            args=args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=True)
-        self.val_transform, _ = build_transform(
-            args=args, trans_config=self.trans_cfg, max_stride=self.model_cfg['max_stride'], is_train=False)
-
-        # ---------------------------- Build Dataset & Dataloader ----------------------------
-        self.dataset, self.dataset_info = build_dataset(args, self.data_cfg, self.trans_cfg, self.train_transform, is_train=True)
-        self.train_loader = build_dataloader(args, self.dataset, self.args.batch_size // self.world_size, CollateFunc())
-
-        # ---------------------------- Build Evaluator ----------------------------
-        self.evaluator = build_evluator(args, self.data_cfg, self.val_transform, self.device)
-
-        # ---------------------------- Build Grad. Scaler ----------------------------
-        self.scaler = torch.cuda.amp.GradScaler(enabled=self.args.fp16)
-
-        # ---------------------------- Build Optimizer ----------------------------
-        self.optimizer_dict['lr0'] *= args.batch_size * self.grad_accumulate / 64
-        self.optimizer, self.start_epoch = build_yolo_optimizer(self.optimizer_dict, model, args.resume)
-
-        # ---------------------------- Build LR Scheduler ----------------------------
-        self.lr_scheduler, self.lf = build_lambda_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
-        if self.args.resume and self.args.resume != 'None':
-            self.lr_scheduler.step()
-
-        # ---------------------------- Build Model-EMA ----------------------------
-        if self.args.ema and distributed_utils.get_rank() in [-1, 0]:
-            print('Build ModelEMA ...')
-            self.model_ema = ModelEMA(self.ema_dict, model, self.start_epoch * len(self.train_loader))
-        else:
-            self.model_ema = None
-
-    def train(self, model):
-        for epoch in range(self.start_epoch, self.args.max_epoch):
-            if self.args.distributed:
-                self.train_loader.batch_sampler.sampler.set_epoch(epoch)
-
-            # check second stage
-            if epoch >= (self.args.max_epoch - self.second_stage_epoch - 1) and not self.second_stage:
-                self.check_second_stage()
-                # save model of the last mosaic epoch
-                weight_name = '{}_last_mosaic_epoch.pth'.format(self.args.model)
-                checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                print('Saving state of the last Mosaic 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)
-
-            # 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)
-
-            # eval one epoch
-            if self.heavy_eval:
-                model_eval = model.module if self.args.distributed else model
-                self.eval(model_eval)
-            else:
-                model_eval = model.module if self.args.distributed else model
-                if (epoch % self.args.eval_epoch) == 0 or (epoch == self.args.max_epoch - 1):
-                    self.eval(model_eval)
-
-            if self.args.debug:
-                print("For debug mode, we only train 1 epoch")
-                break
-
-    def eval(self, model):
-        # chech model
-        model_eval = model if self.model_ema is None else self.model_ema.ema
-
-        if distributed_utils.is_main_process():
-            # check evaluator
-            if self.evaluator is None:
-                print('No evaluator ... save model and go on training.')
-                print('Saving state, epoch: {}'.format(self.epoch))
-                weight_name = '{}_no_eval.pth'.format(self.args.model)
-                checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                torch.save({'model': model_eval.state_dict(),
-                            'mAP': -1.,
-                            'optimizer': self.optimizer.state_dict(),
-                            'epoch': self.epoch,
-                            'args': self.args}, 
-                            checkpoint_path)               
-            else:
-                print('eval ...')
-                # set eval mode
-                model_eval.trainable = False
-                model_eval.eval()
-
-                # evaluate
-                with torch.no_grad():
-                    self.evaluator.evaluate(model_eval)
-
-                # save model
-                cur_map = self.evaluator.map
-                if cur_map > self.best_map:
-                    # update best-map
-                    self.best_map = cur_map
-                    # save model
-                    print('Saving state, epoch:', self.epoch)
-                    weight_name = '{}_best.pth'.format(self.args.model)
-                    checkpoint_path = os.path.join(self.path_to_save, weight_name)
-                    torch.save({'model': model_eval.state_dict(),
-                                'mAP': round(self.best_map*100, 1),
-                                'optimizer': self.optimizer.state_dict(),
-                                'epoch': self.epoch,
-                                'args': self.args}, 
-                                checkpoint_path)                      
-
-                # set train mode.
-                model_eval.trainable = True
-                model_eval.train()
-
-        if self.args.distributed:
-            # wait for all processes to synchronize
-            dist.barrier()
-
-    def train_one_epoch(self, model):
-        metric_logger = MetricLogger(delimiter="  ")
-        metric_logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))
-        metric_logger.add_meter('size', SmoothedValue(window_size=1, fmt='{value:d}'))
-        header = 'Epoch: [{} / {}]'.format(self.epoch, self.args.max_epoch)
-        epoch_size = len(self.train_loader)
-        print_freq = 10
-
-        # basic parameters
-        epoch_size = len(self.train_loader)
-        img_size = self.args.img_size
-        nw = epoch_size * self.args.wp_epoch
-
-        # Train one epoch
-        for iter_i, (images, targets) in enumerate(metric_logger.log_every(self.train_loader, print_freq, header)):
-            ni = iter_i + self.epoch * epoch_size
-            # Warmup
-            if ni <= nw:
-                xi = [0, nw]  # x interp
-                for j, x in enumerate(self.optimizer.param_groups):
-                    # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
-                    x['lr'] = np.interp(
-                        ni, xi, [self.warmup_dict['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * self.lf(self.epoch)])
-                    if 'momentum' in x:
-                        x['momentum'] = np.interp(ni, xi, [self.warmup_dict['warmup_momentum'], self.optimizer_dict['momentum']])
-                                
-            # To device
-            images = images.to(self.device, non_blocking=True).float()
-
-            # Multi scale
-            if self.args.multi_scale:
-                images, targets, img_size = self.rescale_image_targets(
-                    images, targets, self.model_cfg['stride'], self.args.min_box_size, self.model_cfg['multi_scale'])
-            else:
-                targets = self.refine_targets(targets, self.args.min_box_size)
-                
-            # Visualize train targets
-            if self.args.vis_tgt:
-                vis_data(images*255, targets, self.data_cfg['num_classes'])
-
-            # Inference
-            with torch.cuda.amp.autocast(enabled=self.args.fp16):
-                outputs = model(images)
-                # Compute loss
-                loss_dict = self.criterion(outputs=outputs, targets=targets, epoch=self.epoch, task='det_seg_pos')
-                det_loss_dict = loss_dict['det_loss_dict']
-                seg_loss_dict = loss_dict['seg_loss_dict']
-                pos_loss_dict = loss_dict['pos_loss_dict']
-                
-                # TODO: finish the backward + optimize
-
-            # # Update log
-            # metric_logger.update(**loss_dict_reduced)
-            # metric_logger.update(lr=self.optimizer.param_groups[2]["lr"])
-            # metric_logger.update(grad_norm=grad_norm)
-            # metric_logger.update(size=img_size)
-
-            if self.args.debug:
-                print("For debug mode, we only train 1 iteration")
-                break
-
-        # LR Schedule
-        if not self.second_stage:
-            self.lr_scheduler.step()
-
-        # Gather the stats from all processes
-        metric_logger.synchronize_between_processes()
-        print("Averaged stats:", metric_logger)
-
-    def refine_targets(self, targets, min_box_size):
-        # rescale targets
-        for tgt in targets:
-            boxes = tgt["boxes"].clone()
-            labels = tgt["labels"].clone()
-            # refine tgt
-            tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
-            min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
-            keep = (min_tgt_size >= min_box_size)
-
-            tgt["boxes"] = boxes[keep]
-            tgt["labels"] = labels[keep]
-        
-        return targets
-
-    def rescale_image_targets(self, images, targets, stride, min_box_size, multi_scale_range=[0.5, 1.5]):
-        """
-            Deployed for Multi scale trick.
-        """
-        if isinstance(stride, int):
-            max_stride = stride
-        elif isinstance(stride, list):
-            max_stride = max(stride)
-
-        # During training phase, the shape of input image is square.
-        old_img_size = images.shape[-1]
-        new_img_size = random.randrange(old_img_size * multi_scale_range[0], old_img_size * multi_scale_range[1] + max_stride)
-        new_img_size = new_img_size // max_stride * max_stride  # size
-        if new_img_size / old_img_size != 1:
-            # interpolate
-            images = torch.nn.functional.interpolate(
-                                input=images, 
-                                size=new_img_size, 
-                                mode='bilinear', 
-                                align_corners=False)
-        # rescale targets
-        for tgt in targets:
-            boxes = tgt["boxes"].clone()
-            labels = tgt["labels"].clone()
-            boxes = torch.clamp(boxes, 0, old_img_size)
-            # rescale box
-            boxes[:, [0, 2]] = boxes[:, [0, 2]] / old_img_size * new_img_size
-            boxes[:, [1, 3]] = boxes[:, [1, 3]] / old_img_size * new_img_size
-            # refine tgt
-            tgt_boxes_wh = boxes[..., 2:] - boxes[..., :2]
-            min_tgt_size = torch.min(tgt_boxes_wh, dim=-1)[0]
-            keep = (min_tgt_size >= min_box_size)
-
-            tgt["boxes"] = boxes[keep]
-            tgt["labels"] = labels[keep]
-
-        return images, targets, new_img_size
-
-    def check_second_stage(self):
-        # set second stage
-        print('============== Second stage of Training ==============')
-        self.second_stage = True
-
-        # close mosaic augmentation
-        if self.train_loader.dataset.mosaic_prob > 0.:
-            print(' - Close < Mosaic Augmentation > ...')
-            self.train_loader.dataset.mosaic_prob = 0.
-            self.heavy_eval = True
-
-        # close mixup augmentation
-        if self.train_loader.dataset.mixup_prob > 0.:
-            print(' - Close < Mixup Augmentation > ...')
-            self.train_loader.dataset.mixup_prob = 0.
-            self.heavy_eval = True
-
-        # close rotation augmentation
-        if 'degrees' in self.trans_cfg.keys() and self.trans_cfg['degrees'] > 0.0:
-            print(' - Close < degress of rotation > ...')
-            self.trans_cfg['degrees'] = 0.0
-        if 'shear' in self.trans_cfg.keys() and self.trans_cfg['shear'] > 0.0:
-            print(' - Close < shear of rotation >...')
-            self.trans_cfg['shear'] = 0.0
-        if 'perspective' in self.trans_cfg.keys() and self.trans_cfg['perspective'] > 0.0:
-            print(' - Close < perspective of rotation > ...')
-            self.trans_cfg['perspective'] = 0.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
-        
-    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
-   
-
 # Build Trainer
 def build_trainer(args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size):
     # ----------------------- Det trainers -----------------------
@@ -2299,14 +1612,6 @@ def build_trainer(args, data_cfg, model_cfg, trans_cfg, device, model, criterion
     elif model_cfg['trainer_type'] == 'rtpdetr':
         return RTPDetrTrainer(args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size)
     
-    # ----------------------- Det + Seg trainers -----------------------
-    elif model_cfg['trainer_type'] == 'rtcdet_ds':
-        return RTCTrainerDS(args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size)
-
-    # ----------------------- Det + Seg + Pos trainers -----------------------
-    elif model_cfg['trainer_type'] == 'rtcdet_dsp':
-        return RTCTrainerDSP(args, data_cfg, model_cfg, trans_cfg, device, model, criterion, world_size)
-
     else:
         raise NotImplementedError(model_cfg['trainer_type'])
     

models/detectors/rtdetr/basic_modules/basic.py

@@ -254,17 +254,21 @@ class RTCBlock(nn.Module):
                  norm_type  = 'BN',
                  depthwise  = False,):
         super(RTCBlock, self).__init__()
-        self.inter_dim = out_dim
+        self.inter_dim = out_dim // 2
         self.conv1 = BasicConv(in_dim, self.inter_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
         self.conv2 = BasicConv(in_dim, self.inter_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
-        self.m = nn.Sequential(*[Bottleneck(self.inter_dim, self.inter_dim,
-                                            1.0, [1, 3], shortcut,
-                                            act_type, norm_type, depthwise)
-                                            for _ in range(num_blocks)])
-        self.conv3 = BasicConv(self.inter_dim, out_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
+        self.cmodules = nn.ModuleList([Bottleneck(self.inter_dim, self.inter_dim,
+                                                   1.0, [3, 3], shortcut,
+                                                   act_type, norm_type, depthwise)
+                                                   for _ in range(num_blocks)])
+        self.conv3 = BasicConv(self.inter_dim * (2 + num_blocks), out_dim, kernel_size=1, act_type=act_type, norm_type=norm_type)
 
     def forward(self, x):
-        x1 = self.conv1(x)
-        x2 = self.m(self.conv2(x))
-        return self.conv3(x1 + x2)
+        x1, x2 = self.conv1(x), self.conv2(x)
+        out = [x1, x2]
+        for m in self.cmodules:
+            x2 = m(x2)
+            out.append(x2)
+
+        return self.conv3(torch.cat(out, dim=1))
     

utils/solver/optimizer.py

@@ -79,9 +79,10 @@ def build_rtdetr_optimizer(cfg, model, resume=None):
     optimizer.add_param_group({"params": param_dicts[2], "weight_decay": cfg['weight_decay']})
 
     # Backbone's learnable parameters
-    optimizer.add_param_group({"params": param_dicts[3], "lr": cfg['lr0'] * cfg['backbone_lr_ratio'], "weight_decay": 0.0})
-    optimizer.add_param_group({"params": param_dicts[4], "lr": cfg['lr0'] * cfg['backbone_lr_ratio'], "weight_decay": 0.0})
-    optimizer.add_param_group({"params": param_dicts[5], "lr": cfg['lr0'] * cfg['backbone_lr_ratio'], "weight_decay": cfg['weight_decay']})
+    backbone_lr = cfg['lr0'] * cfg['backbone_lr_ratio']
+    optimizer.add_param_group({"params": param_dicts[3], "lr": backbone_lr, "weight_decay": 0.0})
+    optimizer.add_param_group({"params": param_dicts[4], "lr": backbone_lr, "weight_decay": 0.0})
+    optimizer.add_param_group({"params": param_dicts[5], "lr": backbone_lr, "weight_decay": cfg['weight_decay']})
 
     start_epoch = 0
     if resume and resume != 'None':