train RTCDet-S on COCO

models/detectors/rtcdet/loss.py

@@ -302,7 +302,7 @@ class Criterion(object):
         box_targets_pos = box_targets[pos_inds]
         box_weight_pos = assign_metrics[pos_inds]
         loss_box = self.loss_bboxes(box_preds_pos, box_targets_pos)
-        # loss_box *= box_weight_pos
+        loss_box *= box_weight_pos
         loss_box = loss_box.sum() / normalizer
 
         # ------------------ Distribution focal loss  ------------------
@@ -318,7 +318,7 @@ class Criterion(object):
         strides_pos = strides[pos_inds]
         ## compute dfl
         loss_dfl = self.loss_dfl(reg_preds_pos, box_targets_pos, anchors_pos, strides_pos)
-        # loss_dfl *= box_weight_pos
+        loss_dfl *= box_weight_pos
         loss_dfl = loss_dfl.sum() / normalizer
 
         # total loss

test.py

@@ -31,7 +31,7 @@ def parse_args():
                         help='use cuda.')
     parser.add_argument('--save_folder', default='det_results/', type=str,
                         help='Dir to save results')
-    parser.add_argument('-vt', '--visual_threshold', default=0.4, type=float,
+    parser.add_argument('-vt', '--visual_threshold', default=0.3, type=float,
                         help='Final confidence threshold')
     parser.add_argument('-ws', '--window_scale', default=1.0, type=float,
                         help='resize window of cv2 for visualization.')
@@ -125,7 +125,7 @@ def test(args,
     save_path = os.path.join('det_results/', args.dataset, args.model)
     os.makedirs(save_path, exist_ok=True)
 
-    for index in range(num_images):
+    for index in range(2000,num_images):
         print('Testing image {:d}/{:d}....'.format(index+1, num_images))
         image, _ = dataset.pull_image(index)
 

train_multi_gpus.sh

@@ -4,7 +4,7 @@ python -m torch.distributed.run --nproc_per_node=8 train.py \
                                                     -dist \
                                                     -d coco \
                                                     --root /data/datasets/ \
-                                                    -m rtcdet_m \
+                                                    -m rtcdet_s \
                                                     -bs 128 \
                                                     -size 640 \
                                                     --wp_epoch 3 \