yjh0410 il y a 2 ans
Parent
commit
524b08164a
5 fichiers modifiés avec 9 ajouts et 6 suppressions
  1. 3 2
      config/model_config/rtrdet_config.py
  2. 1 1
      engine.py
  3. 3 2
      models/detectors/rtrdet/rtrdet.py
  4. 1 1
      train.py
  5. 1 0
      train_single_gpu.sh

+ 3 - 2
config/model_config/rtrdet_config.py

@@ -12,7 +12,8 @@ rtrdet_cfg = {
         'bk_depthwise': False,
         'width': 1.0,
         'depth': 1.0,
-        'max_stride': 16,
+        'max_stride': 32,
+        'out_stride': 16,
         'd_model': 512,
         ## Transformer Encoder
         'transformer': 'RTRDetTransformer',
@@ -29,7 +30,7 @@ rtrdet_cfg = {
         'decoder_mlp_ratio': 4.0,
         'decoder_dropout': 0.1,
         'decoder_act': 'relu',
-        'decoder_num_queries': 300,
+        'decoder_num_queries': 100,
         'decoder_num_pattern': 3,
         'spatial_prior': 'learned',  # 'learned', 'grid'
         'num_topk': 100,

+ 1 - 1
engine.py

@@ -1287,7 +1287,7 @@ class RTRTrainer(object):
             # Multi scale
             if self.args.multi_scale:
                 images, targets, img_size = self.rescale_image_targets(
-                    images, targets, 32, self.args.min_box_size, self.model_cfg['multi_scale'])
+                    images, targets, self.model_cfg['max_stride'], self.args.min_box_size, self.model_cfg['multi_scale'])
             else:
                 targets = self.refine_targets(targets, self.args.min_box_size)
 

+ 3 - 2
models/detectors/rtrdet/rtrdet.py

@@ -15,12 +15,13 @@ class RTRDet(nn.Module):
                  aux_loss    :bool = False,
                  deploy      :bool = False):
         super(RTRDet, self).__init__()
-        assert cfg['max_stride'] == 16 or cfg['max_stride'] == 32
+        assert cfg['out_stride'] == 16 or cfg['out_stride'] == 32
         # ------------------ Basic parameters ------------------
         self.cfg = cfg
         self.device = device
+        self.out_stride = cfg['out_stride']
         self.max_stride = cfg['max_stride']
-        self.num_levels = 2 if cfg['max_stride'] == 16 else 1
+        self.num_levels = 2 if cfg['out_stride'] == 16 else 1
         self.num_topk = cfg['num_topk']
         self.num_classes = num_classes
         self.d_model = round(cfg['d_model'] * cfg['width'])

+ 1 - 1
train.py

@@ -23,7 +23,7 @@ from engine import build_trainer
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='YOLO-Tutorial')
+    parser = argparse.ArgumentParser(description='Real-time Object Detection LAB')
     # Basic
     parser.add_argument('--cuda', action='store_true', default=False,
                         help='use cuda.')

+ 1 - 0
train_single_gpu.sh

@@ -14,6 +14,7 @@ python train.py \
         --ema \
         --fp16 \
         --multi_scale \
+        --eval_first \
         # --load_cache \
         # --resume weights/coco/yolox_m/yolox_m_best.pth \
         # --eval_first