Add train_cfg/test_cfg in AnchorHead (open-mmlab#2422)

* add cfg for anchor head

* add copy in two stage rpn

Author: xvjiarui

mmdet/models/anchor_heads/anchor_head.py

@@ -49,7 +49,9 @@ def __init__(self,
                      use_sigmoid=True,
                      loss_weight=1.0),
                  loss_bbox=dict(
-                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)):
+                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0),
+                 train_cfg=None,
+                 test_cfg=None):
         super(AnchorHead, self).__init__()
         self.in_channels = in_channels
         self.num_classes = num_classes
@@ -80,6 +82,8 @@ def __init__(self,
 
         self.loss_cls = build_loss(loss_cls)
         self.loss_bbox = build_loss(loss_bbox)
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
         self.fp16_enabled = False
 
         self.anchor_generators = []
@@ -149,7 +153,7 @@ def get_anchors(self, featmap_sizes, img_metas, device='cuda'):
         return anchor_list, valid_flag_list
 
     def loss_single(self, cls_score, bbox_pred, labels, label_weights,
-                    bbox_targets, bbox_weights, num_total_samples, cfg):
+                    bbox_targets, bbox_weights, num_total_samples):
         # classification loss
         labels = labels.reshape(-1)
         label_weights = label_weights.reshape(-1)
@@ -175,7 +179,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
         assert len(featmap_sizes) == len(self.anchor_generators)
@@ -192,7 +195,7 @@ def loss(self,
             img_metas,
             self.target_means,
             self.target_stds,
-            cfg,
+            self.train_cfg,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             label_channels=label_channels,
@@ -212,8 +215,7 @@ def loss(self,
             label_weights_list,
             bbox_targets_list,
             bbox_weights_list,
-            num_total_samples=num_total_samples,
-            cfg=cfg)
+            num_total_samples=num_total_samples)
         return dict(loss_cls=losses_cls, loss_bbox=losses_bbox)
 
     @force_fp32(apply_to=('cls_scores', 'bbox_preds'))

mmdet/models/anchor_heads/atss_head.py

@@ -123,7 +123,7 @@ def forward_single(self, x, scale):
         return cls_score, bbox_pred, centerness
 
     def loss_single(self, anchors, cls_score, bbox_pred, centerness, labels,
-                    label_weights, bbox_targets, num_total_samples, cfg):
+                    label_weights, bbox_targets, num_total_samples):
 
         anchors = anchors.reshape(-1, 4)
         cls_score = cls_score.permute(0, 2, 3,
@@ -186,7 +186,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
 
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
@@ -202,7 +201,6 @@ def loss(self,
             valid_flag_list,
             gt_bboxes,
             img_metas,
-            cfg,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             label_channels=label_channels)
@@ -226,8 +224,7 @@ def loss(self,
                 labels_list,
                 label_weights_list,
                 bbox_targets_list,
-                num_total_samples=num_total_samples,
-                cfg=cfg)
+                num_total_samples=num_total_samples)
 
         bbox_avg_factor = sum(bbox_avg_factor)
         bbox_avg_factor = reduce_mean(bbox_avg_factor).item()

mmdet/models/anchor_heads/fcos_head.py

@@ -53,7 +53,9 @@ def __init__(self,
                      use_sigmoid=True,
                      loss_weight=1.0),
                  conv_cfg=None,
-                 norm_cfg=dict(type='GN', num_groups=32, requires_grad=True)):
+                 norm_cfg=dict(type='GN', num_groups=32, requires_grad=True),
+                 train_cfg=None,
+                 test_cfg=None):
         super(FCOSHead, self).__init__()
         self.num_classes = num_classes
         self.cls_out_channels = num_classes
@@ -65,6 +67,8 @@ def __init__(self,
         self.loss_cls = build_loss(loss_cls)
         self.loss_bbox = build_loss(loss_bbox)
         self.loss_centerness = build_loss(loss_centerness)
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
         self.conv_cfg = conv_cfg
         self.norm_cfg = norm_cfg
         self.fp16_enabled = False
@@ -147,7 +151,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         assert len(cls_scores) == len(bbox_preds) == len(centernesses)
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]

mmdet/models/anchor_heads/fovea_head.py

@@ -62,7 +62,9 @@ def __init__(self,
                  loss_cls=None,
                  loss_bbox=None,
                  conv_cfg=None,
-                 norm_cfg=None):
+                 norm_cfg=None,
+                 train_cfg=None,
+                 test_cfg=None):
         super(FoveaHead, self).__init__()
         self.num_classes = num_classes
         self.cls_out_channels = num_classes
@@ -84,6 +86,8 @@ def __init__(self,
         self.loss_bbox = build_loss(loss_bbox)
         self.conv_cfg = conv_cfg
         self.norm_cfg = norm_cfg
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
         self._init_layers()
 
     def _init_layers(self):
@@ -195,7 +199,6 @@ def loss(self,
              gt_bbox_list,
              gt_label_list,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         assert len(cls_scores) == len(bbox_preds)
 

mmdet/models/anchor_heads/free_anchor_retina_head.py

@@ -38,7 +38,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
         assert len(featmap_sizes) == len(self.anchor_generators)

mmdet/models/anchor_heads/ga_rpn_head.py

@@ -40,7 +40,6 @@ def loss(self,
              loc_preds,
              gt_bboxes,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         losses = super(GARPNHead, self).loss(
             cls_scores,
@@ -50,7 +49,6 @@ def loss(self,
             gt_bboxes,
             None,
             img_metas,
-            cfg,
             gt_bboxes_ignore=gt_bboxes_ignore)
         return dict(
             loss_rpn_cls=losses['loss_cls'],

mmdet/models/anchor_heads/guided_anchor_head.py

@@ -113,6 +113,8 @@ def __init__(
         deformable_groups=4,
         loc_filter_thr=0.01,
         background_label=None,
+        train_cfg=None,
+        test_cfg=None,
         loss_loc=dict(
             type='FocalLoss',
             use_sigmoid=True,
@@ -176,6 +178,9 @@ def __init__(
         self.loss_cls = build_loss(loss_cls)
         self.loss_bbox = build_loss(loss_bbox)
 
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
+
         self.fp16_enabled = False
 
         self._init_layers()
@@ -418,7 +423,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
         assert len(featmap_sizes) == len(self.approx_generators)
@@ -431,26 +435,26 @@ def loss(self,
             featmap_sizes,
             self.octave_base_scale,
             self.anchor_strides,
-            center_ratio=cfg.center_ratio,
-            ignore_ratio=cfg.ignore_ratio)
+            center_ratio=self.train_cfg.center_ratio,
+            ignore_ratio=self.train_cfg.ignore_ratio)
 
         # get sampled approxes
         approxs_list, inside_flag_list = self.get_sampled_approxs(
-            featmap_sizes, img_metas, cfg, device=device)
+            featmap_sizes, img_metas, self.train_cfg, device=device)
         # get squares and guided anchors
         squares_list, guided_anchors_list, _ = self.get_anchors(
             featmap_sizes, shape_preds, loc_preds, img_metas, device=device)
 
         # get shape targets
-        sampling = False if not hasattr(cfg, 'ga_sampler') else True
+        sampling = False if not hasattr(self.train_cfg, 'ga_sampler') else True
         shape_targets = ga_shape_target(
             approxs_list,
             inside_flag_list,
             squares_list,
             gt_bboxes,
             img_metas,
             self.approxs_per_octave,
-            cfg,
+            self.train_cfg,
             sampling=sampling)
         if shape_targets is None:
             return None
@@ -469,7 +473,7 @@ def loss(self,
             img_metas,
             self.target_means,
             self.target_stds,
-            cfg,
+            self.train_cfg,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             label_channels=label_channels,
@@ -492,8 +496,7 @@ def loss(self,
             label_weights_list,
             bbox_targets_list,
             bbox_weights_list,
-            num_total_samples=num_total_samples,
-            cfg=cfg)
+            num_total_samples=num_total_samples)
 
         # get anchor location loss
         losses_loc = []
@@ -503,7 +506,7 @@ def loss(self,
                 loc_targets[i],
                 loc_weights[i],
                 loc_avg_factor=loc_avg_factor,
-                cfg=cfg)
+                cfg=self.train_cfg)
             losses_loc.append(loss_loc)
 
         # get anchor shape loss

mmdet/models/anchor_heads/reppoints_head.py

@@ -64,7 +64,9 @@ def __init__(self,
                  use_grid_points=False,
                  center_init=True,
                  transform_method='moment',
-                 moment_mul=0.01):
+                 moment_mul=0.01,
+                 train_cfg=None,
+                 test_cfg=None):
         super(RepPointsHead, self).__init__()
         self.in_channels = in_channels
         self.num_classes = num_classes
@@ -77,6 +79,8 @@ def __init__(self,
         self.point_strides = point_strides
         self.conv_cfg = conv_cfg
         self.norm_cfg = norm_cfg
+        self.train_cfg = train_cfg
+        self.test_cfg = test_cfg
 
         self.background_label = (
             num_classes if background_label is None else background_label)
@@ -423,7 +427,6 @@ def loss(self,
              gt_bboxes,
              gt_labels,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
         assert len(featmap_sizes) == len(self.point_generators)
@@ -434,7 +437,7 @@ def loss(self,
                                                        img_metas)
         pts_coordinate_preds_init = self.offset_to_pts(center_list,
                                                        pts_preds_init)
-        if cfg.init.assigner['type'] == 'PointAssigner':
+        if self.train_cfg.init.assigner['type'] == 'PointAssigner':
             # Assign target for center list
             candidate_list = center_list
         else:
@@ -447,7 +450,7 @@ def loss(self,
             valid_flag_list,
             gt_bboxes,
             img_metas,
-            cfg.init,
+            self.train_cfg.init,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             label_channels=label_channels,
@@ -481,7 +484,7 @@ def loss(self,
             valid_flag_list,
             gt_bboxes,
             img_metas,
-            cfg.refine,
+            self.train_cfg.refine,
             gt_bboxes_ignore_list=gt_bboxes_ignore,
             gt_labels_list=gt_labels,
             label_channels=label_channels,

mmdet/models/anchor_heads/rpn_head.py

@@ -40,15 +40,13 @@ def loss(self,
              bbox_preds,
              gt_bboxes,
              img_metas,
-             cfg,
              gt_bboxes_ignore=None):
         losses = super(RPNHead, self).loss(
             cls_scores,
             bbox_preds,
             gt_bboxes,
             None,
             img_metas,
-            cfg,
             gt_bboxes_ignore=gt_bboxes_ignore)
         return dict(
             loss_rpn_cls=losses['loss_cls'], loss_rpn_bbox=losses['loss_bbox'])