Add fused Mask Traget generator op to Mask-RCNN

Signed-off-by: Serge Panev <[email protected]>
dmlc · Oct 17, 2019 · 7af2157 · 7af2157
1 parent 31f91d2
commit 7af2157
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/gluoncv/model_zoo/mask_rcnn/rcnn_target.py b/gluoncv/model_zoo/mask_rcnn/rcnn_target.py
@@ -47,6 +47,12 @@ def hybrid_forward(self, F, rois, gt_masks, matches, cls_targets):
 
         """
 
+        if hasattr(F.contrib, 'mrcnn_mask_target'):
+            return F.contrib.mrcnn_mask_target(rois, gt_masks, matches, cls_targets,
+                                               num_rois=self._num_rois,
+                                               num_classes=self._num_classes,
+                                               mask_size=self._mask_size)
+
         # cannot know M (num_gt) to have accurate batch id B * M, must split batch dim
         def _split(x, axis, num_outputs, squeeze_axis):
             x = F.split(x, axis=axis, num_outputs=num_outputs, squeeze_axis=squeeze_axis)