Fix Faster-RCNN d2_convert error. (#1389)

tensorpack · Jan 22, 2020 · 12ad257 · 12ad257
1 parent 882f858
commit 12ad257
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/examples/FasterRCNN/convert_d2/convert_d2.py b/examples/FasterRCNN/convert_d2/convert_d2.py
@@ -109,10 +109,11 @@ def _convert_box_predictor(src, dst):
         _convert_box_predictor("roi_heads.box_predictor", "fastrcnn/outputs" if has_fpn else "fastrcnn")
 
     # mask head
-    for fcn in range(cfg.MODEL.ROI_MASK_HEAD.NUM_CONV):
-        _convert_conv(f"roi_heads.mask_head.mask_fcn{fcn+1}", f"maskrcnn/fcn{fcn}")
-    _convert_conv("roi_heads.mask_head.deconv", "maskrcnn/deconv")
-    _convert_conv("roi_heads.mask_head.predictor", "maskrcnn/conv")
+    if cfg.MODEL.MASK_ON:
+        for fcn in range(cfg.MODEL.ROI_MASK_HEAD.NUM_CONV):
+            _convert_conv(f"roi_heads.mask_head.mask_fcn{fcn+1}", f"maskrcnn/fcn{fcn}")
+        _convert_conv("roi_heads.mask_head.deconv", "maskrcnn/deconv")
+        _convert_conv("roi_heads.mask_head.predictor", "maskrcnn/conv")
 
     for k in list(d.keys()):
         if "cell_anchors" in k: