feat(detection): support Objects365 and reformat

Author: wangjianfeng

official/vision/detection/layers/basic/functional.py

@@ -10,8 +10,7 @@
 import megengine.functional as F
 import numpy as np
 
-from megengine import _internal as mgb
-from megengine.core import Tensor, wrap_io_tensor
+from megengine.core import Tensor
 
 
 def get_padded_array_np(
@@ -86,8 +85,3 @@ def get_padded_tensor(
     else:
         raise Exception("Not supported tensor dim: %d" % ndim)
     return padded_array
-
-
-@wrap_io_tensor
-def indexing_set_one_hot(inp, axis, idx, value) -> Tensor:
-    return mgb.opr.indexing_set_one_hot(inp, axis, idx, value)

official/vision/detection/layers/det/loss.py

@@ -12,8 +12,6 @@
 
 from megengine.core import tensor, Tensor
 
-from official.vision.detection.layers import basic
-
 
 def get_focal_loss(
     score: Tensor,
@@ -51,28 +49,19 @@ def get_focal_loss(
     Returns:
         the calculated focal loss.
     """
-    mask = 1 - (label == ignore_label)
-    valid_label = label * mask
-
-    score_shp = score.shape
-    zero_mat = mge.zeros(
-        F.concat([score_shp[0], score_shp[1], score_shp[2] + 1], axis=0),
-        dtype=np.float32,
-    )
-    one_mat = mge.ones(
-        F.concat([score_shp[0], score_shp[1], tensor(1)], axis=0), dtype=np.float32,
-    )
-
-    one_hot = basic.indexing_set_one_hot(
-        zero_mat, 2, valid_label.astype(np.int32), one_mat
-    )[:, :, 1:]
-    pos_part = F.power(1 - score, gamma) * one_hot * F.log(score)
-    neg_part = F.power(score, gamma) * (1 - one_hot) * F.log(1 - score)
-    loss = -(alpha * pos_part + (1 - alpha) * neg_part).sum(axis=2) * mask
+    class_range = F.arange(1, score.shape[2] + 1)
+
+    label = F.add_axis(label, axis=2)
+    pos_part = (1 - score) ** gamma * F.log(score)
+    neg_part = score ** gamma * F.log(1 - score)
+
+    pos_loss = -(label == class_range) * pos_part * alpha
+    neg_loss = -(label != class_range) * (label != ignore_label) * neg_part * (1 - alpha)
+    loss = pos_loss + neg_loss
 
     if norm_type == "fg":
-        positive_mask = label > background
-        return loss.sum() / F.maximum(positive_mask.sum(), 1)
+        fg_mask = (label != background) * (label != ignore_label)
+        return loss.sum() / F.maximum(fg_mask.sum(), 1)
     elif norm_type == "none":
         return loss.sum()
     else:
@@ -117,8 +106,7 @@ def get_smooth_l1_loss(
     gt_bbox = gt_bbox.reshape(-1, 4)
     label = label.reshape(-1)
 
-    valid_mask = 1 - (label == ignore_label)
-    fg_mask = (1 - (label == background)) * valid_mask
+    fg_mask = (label != background) * (label != ignore_label)
 
     losses = get_smooth_l1_base(pred_bbox, gt_bbox, sigma, is_fix=fix_smooth_l1)
     if norm_type == "fg":
@@ -154,19 +142,16 @@ def get_smooth_l1_base(
         cond_point = sigma
         x = pred_bbox - gt_bbox
         abs_x = F.abs(x)
-        in_mask = abs_x < cond_point
-        out_mask = 1 - in_mask
-        in_loss = 0.5 * (x ** 2)
-        out_loss = sigma * abs_x - 0.5 * (sigma ** 2)
-        loss = in_loss * in_mask + out_loss * out_mask
+        in_loss = 0.5 * x ** 2
+        out_loss = sigma * abs_x - 0.5 * sigma ** 2
     else:
         sigma2 = sigma ** 2
         cond_point = 1 / sigma2
         x = pred_bbox - gt_bbox
         abs_x = F.abs(x)
-        in_mask = abs_x < cond_point
-        out_mask = 1 - in_mask
-        in_loss = 0.5 * (sigma * x) ** 2
+        in_loss = 0.5 * x ** 2 * sigma2
         out_loss = abs_x - 0.5 / sigma2
-        loss = in_loss * in_mask + out_loss * out_mask
+    in_mask = abs_x < cond_point
+    out_mask = 1 - in_mask
+    loss = in_loss * in_mask + out_loss * out_mask
     return loss

official/vision/detection/layers/det/retinanet.py

@@ -28,7 +28,7 @@ def __init__(self, cfg, input_shape: List[basic.ShapeSpec]):
         num_classes = cfg.num_classes
         num_convs = 4
         prior_prob = cfg.cls_prior_prob
-        num_anchors = [9, 9, 9, 9, 9]
+        num_anchors = [len(cfg.anchor_ratios) * len(cfg.anchor_scales)] * 5
 
         assert (
             len(set(num_anchors)) == 1

official/vision/detection/models/__init__.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+from .retinanet import *
+
+_EXCLUDE = {}
+__all__ = [k for k in globals().keys() if k not in _EXCLUDE and not k.startswith("_")]

official/vision/detection/retinanet_res50_1x_800size.py renamed to official/vision/detection/models/retinanet.py

@@ -10,7 +10,6 @@
 import megengine.functional as F
 import megengine.module as M
 import numpy as np
-from megengine import hub
 
 from official.vision.classification.resnet.model import resnet50
 from official.vision.detection import layers
@@ -47,7 +46,7 @@ def __init__(self, cfg, batch_size):
             for p in bottom_up.layer1.parameters():
                 p.requires_grad = False
 
-        # -------------------------- build the FPN -------------------------- #
+        # ----------------------- build the FPN ----------------------------- #
         in_channels_p6p7 = 2048
         out_channels = 256
         self.backbone = layers.FPN(
@@ -61,7 +60,7 @@ def __init__(self, cfg, batch_size):
         backbone_shape = self.backbone.output_shape()
         feature_shapes = [backbone_shape[f] for f in self.in_features]
 
-        # -------------------------- build the RetinaNet Head -------------- #
+        # ----------------------- build the RetinaNet Head ------------------ #
         self.head = layers.RetinaNetHead(cfg, feature_shapes)
 
         self.inputs = {
@@ -199,13 +198,22 @@ def __init__(self):
         self.resnet_norm = "FrozenBN"
         self.backbone_freeze_at = 2
 
-        # ------------------------ data cfg --------------------------- #
+        # ------------------------ data cfg -------------------------- #
+        self.train_dataset = dict(
+            name="coco",
+            root="train2017",
+            ann_file="instances_train2017.json"
+        )
+        self.test_dataset = dict(
+            name="coco",
+            root="val2017",
+            ann_file="instances_val2017.json"
+        )
         self.train_image_short_size = 800
         self.train_image_max_size = 1333
         self.num_classes = 80
         self.img_mean = np.array([103.530, 116.280, 123.675])  # BGR
         self.img_std = np.array([57.375, 57.120, 58.395])
-        # self.img_std = np.array([1.0, 1.0, 1.0])
         self.reg_mean = None
         self.reg_std = np.array([0.1, 0.1, 0.2, 0.2])
 
@@ -217,7 +225,7 @@ def __init__(self):
         self.class_aware_box = False
         self.cls_prior_prob = 0.01
 
-        # ------------------------ losss cfg ------------------------- #
+        # ------------------------ loss cfg -------------------------- #
         self.focal_loss_alpha = 0.25
         self.focal_loss_gamma = 2
         self.reg_loss_weight = 1.0 / 4.0
@@ -229,29 +237,14 @@ def __init__(self):
         self.log_interval = 20
         self.nr_images_epoch = 80000
         self.max_epoch = 18
-        self.warm_iters = 100
+        self.warm_iters = 500
         self.lr_decay_rate = 0.1
         self.lr_decay_sates = [12, 16, 17]
 
-        # ------------------------ testing cfg ------------------------- #
+        # ------------------------ testing cfg ----------------------- #
         self.test_image_short_size = 800
         self.test_image_max_size = 1333
         self.test_max_boxes_per_image = 100
         self.test_vis_threshold = 0.3
         self.test_cls_threshold = 0.05
         self.test_nms = 0.5
-
-
-@hub.pretrained(
-    "https://data.megengine.org.cn/models/weights/"
-    "retinanet_d3f58dce_res50_1x_800size_36dot0.pkl"
-)
-def retinanet_res50_1x_800size(batch_size=1, **kwargs):
-    r"""ResNet-18 model from
-    `"RetinaNet" <https://arxiv.org/abs/1708.02002>`_
-    """
-    return RetinaNet(RetinaNetConfig(), batch_size=batch_size, **kwargs)
-
-
-Net = RetinaNet
-Cfg = RetinaNetConfig

official/vision/detection/retinanet_res50_coco_1x_800size.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+from megengine import hub
+
+from official.vision.detection import models
+
+
+class CustomRetinaNetConfig(models.RetinaNetConfig):
+    def __init__(self):
+        super().__init__()
+
+        # ------------------------ data cfg -------------------------- #
+        self.train_dataset = dict(
+            name="coco",
+            root="train2017",
+            ann_file="annotations/instances_train2017.json"
+        )
+        self.test_dataset = dict(
+            name="coco",
+            root="val2017",
+            ann_file="annotations/instances_val2017.json"
+        )
+
+
+@hub.pretrained(
+    "https://data.megengine.org.cn/models/weights/"
+    "retinanet_d3f58dce_res50_1x_800size_36dot0.pkl"
+)
+def retinanet_res50_coco_1x_800size(batch_size=1, **kwargs):
+    r"""ResNet-18 model from
+    `"RetinaNet" <https://arxiv.org/abs/1708.02002>`_
+    """
+    return models.RetinaNet(RetinaNetConfig(), batch_size=batch_size, **kwargs)
+
+
+Net = models.RetinaNet
+Cfg = CustomRetinaNetConfig

official/vision/detection/retinanet_res50_objects365_1x_800size.py

@@ -0,0 +1,42 @@
+# -*- coding: utf-8 -*-
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+from megengine import hub
+
+from official.vision.detection import models
+
+
+class CustomRetinaNetConfig(models.RetinaNetConfig):
+    def __init__(self):
+        super().__init__()
+
+        # ------------------------ data cfg -------------------------- #
+        self.train_dataset = dict(
+            name="objects365",
+            root="train",
+            ann_file="annotations/objects365_train_20190423.json"
+        )
+        self.test_dataset = dict(
+            name="objects365",
+            root="val",
+            ann_file="annotations/objects365_val_20190423.json"
+        )
+
+        # ------------------------ training cfg ---------------------- #
+        self.nr_images_epoch = 400000
+
+
+def retinanet_objects365_res50_1x_800size(batch_size=1, **kwargs):
+    r"""ResNet-18 model from
+    `"RetinaNet" <https://arxiv.org/abs/1708.02002>`_
+    """
+    return models.RetinaNet(RetinaNetConfig(), batch_size=batch_size, **kwargs)
+
+
+Net = models.RetinaNet
+Cfg = CustomRetinaNetConfig

official/vision/detection/tools/data_mapper.py

@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+from megengine.data.dataset import COCO, Objects365
+
+data_mapper = dict(
+    coco=COCO,
+    objects365=Objects365,
+)