add FCOS to d2go

Summary: D2 (https://github.com/facebookresearch/d2go/commit/7992f91324aee6ae59795063a007c6837e60cdb8) does not add new yacs config for new models, but this simple wrapper with configs->arguments mapping is enough to make the model work with yacs config.

Reviewed By: zhanghang1989

Differential Revision: D30980180

fbshipit-source-id: 75a0cc66051800a3e9d553bb650ca5c900c0ffa3

d2go/modeling/fcos.py

+from d2go.config import CfgNode as CN
+from detectron2.config import configurable
+from detectron2.modeling.backbone import build_backbone
+from detectron2.modeling.meta_arch import META_ARCH_REGISTRY
+from detectron2.modeling.meta_arch.fcos import FCOS as d2_FCOS, FCOSHead
+
+
+def add_fcos_configs(cfg):
+    cfg.MODEL.FCOS = CN()
+    # the number of foreground classes.
+    cfg.MODEL.FCOS.NUM_CLASSES = 80
+    cfg.MODEL.FCOS.IN_FEATURES = ["p3", "p4", "p5", "p6", "p7"]
+    cfg.MODEL.FCOS.NUM_CONVS = 4
+    cfg.MODEL.FCOS.HEAD_NORM = "GN"
+
+    # inference parameters
+    cfg.MODEL.FCOS.SCORE_THRESH_TEST = 0.04
+    cfg.MODEL.FCOS.TOPK_CANDIDATES_TEST = 1000
+    cfg.MODEL.FCOS.NMS_THRESH_TEST = 0.6
+
+    # Focal loss parameters
+    cfg.MODEL.FCOS.FOCAL_LOSS_ALPHA = 0.25
+    cfg.MODEL.FCOS.FOCAL_LOSS_GAMMA = 2.0
+
+
+@META_ARCH_REGISTRY.register()
+class FCOS(d2_FCOS):
+    """
+    Implement config->argument translation for FCOS model.
+    """
+
+    @configurable
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def from_config(cls, cfg):
+        backbone = build_backbone(cfg)
+        backbone_shape = backbone.output_shape()
+        try:
+            feature_shapes = [backbone_shape[f] for f in cfg.MODEL.FCOS.IN_FEATURES]
+        except KeyError:
+            raise KeyError(
+                f"Available keys: {backbone_shape.keys()}.  Requested keys: {cfg.MODEL.FCOS.IN_FEATURES}"
+            )
+        head = FCOSHead(
+            input_shape=feature_shapes,
+            num_classes=cfg.MODEL.FCOS.NUM_CLASSES,
+            conv_dims=[feature_shapes[0].channels] * cfg.MODEL.FCOS.NUM_CONVS,
+            norm=cfg.MODEL.FCOS.HEAD_NORM,
+        )
+        return {
+            "backbone": backbone,
+            "head": head,
+            "pixel_mean": cfg.MODEL.PIXEL_MEAN,
+            "pixel_std": cfg.MODEL.PIXEL_STD,
+            "num_classes": cfg.MODEL.FCOS.NUM_CLASSES,
+            "head_in_features": cfg.MODEL.FCOS.IN_FEATURES,
+            # Loss parameters:
+            "focal_loss_alpha": cfg.MODEL.FCOS.FOCAL_LOSS_ALPHA,
+            "focal_loss_gamma": cfg.MODEL.FCOS.FOCAL_LOSS_GAMMA,
+            # Inference parameters:
+            "test_score_thresh": cfg.MODEL.FCOS.SCORE_THRESH_TEST,
+            "test_topk_candidates": cfg.MODEL.FCOS.TOPK_CANDIDATES_TEST,
+            "test_nms_thresh": cfg.MODEL.FCOS.NMS_THRESH_TEST,
+            "max_detections_per_image": cfg.TEST.DETECTIONS_PER_IMAGE,
+        }

d2go/utils/get_default_cfg.py

@@ -7,10 +7,8 @@ from d2go.data.build import (
 )
 from d2go.data.config import add_d2go_data_default_configs
 from d2go.modeling import kmeans_anchors, model_ema
-from d2go.modeling.backbone.fbnet_cfg import (
-    add_bifpn_default_configs,
-    add_fbnet_v2_default_configs,
-)
+from d2go.modeling.backbone.fbnet_cfg import add_fbnet_v2_default_configs
+from d2go.modeling.fcos import add_fcos_configs
 from d2go.modeling.model_freezing_utils import add_model_freezing_configs
 from d2go.modeling.quantization import add_quantization_default_configs
 from d2go.modeling.subclass import add_subclass_configs
@@ -59,6 +57,8 @@ def get_default_cfg(_C):
     add_abnormal_checker_configs(_C)
     # _C.MODEL.SUBCLASS
     add_subclass_configs(_C)
+    # _C.MODEL.FCOS
+    add_fcos_configs(_C)
 
     # Set find_unused_parameters for DistributedDataParallel.
     _C.MODEL.DDP_FIND_UNUSED_PARAMETERS = False