[Fix]: fix mask rcnn training stuck problem when there is no positive rois (#3713)

* Fix mask rcnn stuck problem when there is no positive rois

* support non pos inference in cascade methods, link CU-49tawu

* print mmcv version in CI

* use mmcv repo to check wrappers

* change cpu build

* upgrade mmcv requirements and change ci back

* use pre-built whl in CI

.github/workflows/build.yml

@@ -49,7 +49,9 @@ jobs:
       - name: Install PyTorch
         run: pip install torch==${{matrix.torch}}+cpu torchvision==${{matrix.torchvision}}+cpu -f https://download.pytorch.org/whl/torch_stable.html
       - name: Install MMCV
-        run: pip install mmcv-full==latest+torch${{matrix.torch}}+cpu -f https://openmmlab.oss-accelerate.aliyuncs.com/mmcv/dist/index.html
+        run: |
+          pip install mmcv-full==latest+torch${{matrix.torch}}+cpu -f https://openmmlab.oss-accelerate.aliyuncs.com/mmcv/dist/index.html
+          python -c 'import mmcv; print(mmcv.__version__)'
       - name: Install unittest dependencies
         run: pip install -r requirements/tests.txt -r requirements/optional.txt
       - name: Build and install
@@ -118,6 +120,7 @@ jobs:
         run: |
           pip install mmcv-full==${{matrix.mmcv}} -f https://openmmlab.oss-accelerate.aliyuncs.com/mmcv/dist/index.html
           pip install -r requirements.txt
+          python -c 'import mmcv; print(mmcv.__version__)'
       - name: Build and install
         run: |
           rm -rf .eggs

docs/install.md

@@ -72,7 +72,7 @@ pip install mmcv-full
 
 | MMDetection version |    MMCV version     |
 |:-------------------:|:-------------------:|
-| master              | mmcv-full>=1.1.1, <=1.2|
+| master              | mmcv-full>=1.1.5, <=1.2|
 | 2.4.0               | mmcv-full>=1.1.1, <=1.2|
 | 2.3.0               | mmcv-full==1.0.5|
 | 2.3.0rc0            | mmcv-full>=1.0.2    |

mmdet/__init__.py

@@ -15,7 +15,7 @@ def digit_version(version_str):
     return digit_version
 
 
-mmcv_minimum_version = '1.1.3'
+mmcv_minimum_version = '1.1.5'
 mmcv_maximum_version = '1.2'
 mmcv_version = digit_version(mmcv.__version__)
 

mmdet/models/roi_heads/bbox_heads/bbox_head.py

@@ -181,7 +181,7 @@ class BBoxHead(nn.Module):
                     avg_factor=bbox_targets.size(0),
                     reduction_override=reduction_override)
             else:
-                losses['loss_bbox'] = bbox_pred.sum() * 0
+                losses['loss_bbox'] = bbox_pred[pos_inds].sum()
         return losses
 
     @force_fp32(apply_to=('cls_score', 'bbox_pred'))

mmdet/models/roi_heads/cascade_roi_head.py

@@ -187,10 +187,6 @@ class CascadeRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
         """Run forward function and calculate loss for mask head in
         training."""
         pos_rois = bbox2roi([res.pos_bboxes for res in sampling_results])
-        if len(pos_rois) == 0:
-            # If there are no predicted and/or truth boxes, then we cannot
-            # compute head / mask losses
-            return dict(loss_mask=None)
         mask_results = self._mask_forward(stage, x, pos_rois)
 
         mask_targets = self.mask_head[stage].get_targets(
@@ -271,11 +267,9 @@ class CascadeRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
                 mask_results = self._mask_forward_train(
                     i, x, sampling_results, gt_masks, rcnn_train_cfg,
                     bbox_results['bbox_feats'])
-                # TODO: Support empty tensor input. #2280
-                if mask_results['loss_mask'] is not None:
-                    for name, value in mask_results['loss_mask'].items():
-                        losses[f's{i}.{name}'] = (
-                            value * lw if 'loss' in name else value)
+                for name, value in mask_results['loss_mask'].items():
+                    losses[f's{i}.{name}'] = (
+                        value * lw if 'loss' in name else value)
 
             # refine bboxes
             if i < self.num_stages - 1:

mmdet/models/roi_heads/dynamic_roi_head.py

@@ -94,9 +94,7 @@ class DynamicRoIHead(StandardRoIHead):
             mask_results = self._mask_forward_train(x, sampling_results,
                                                     bbox_results['bbox_feats'],
                                                     gt_masks, img_metas)
-            # TODO: Support empty tensor input. #2280
-            if mask_results['loss_mask'] is not None:
-                losses.update(mask_results['loss_mask'])
+            losses.update(mask_results['loss_mask'])
 
         # update IoU threshold and SmoothL1 beta
         update_iter_interval = self.train_cfg.dynamic_rcnn.update_iter_interval

mmdet/models/roi_heads/mask_heads/coarse_mask_head.py

@@ -1,5 +1,5 @@
 import torch.nn as nn
-from mmcv.cnn import ConvModule, constant_init, xavier_init
+from mmcv.cnn import ConvModule, Linear, constant_init, xavier_init
 from mmcv.runner import auto_fp16
 
 from mmdet.models.builder import HEADS
@@ -64,10 +64,10 @@ class CoarseMaskHead(FCNMaskHead):
         for i in range(num_fcs):
             fc_in_channels = (
                 last_layer_dim if i == 0 else self.fc_out_channels)
-            self.fcs.append(nn.Linear(fc_in_channels, self.fc_out_channels))
+            self.fcs.append(Linear(fc_in_channels, self.fc_out_channels))
         last_layer_dim = self.fc_out_channels
         output_channels = self.num_classes * self.output_area
-        self.fc_logits = nn.Linear(last_layer_dim, output_channels)
+        self.fc_logits = Linear(last_layer_dim, output_channels)
 
     def init_weights(self):
         for m in self.fcs.modules():

mmdet/models/roi_heads/mask_heads/fcn_mask_head.py

@@ -138,7 +138,7 @@ class FCNMaskHead(nn.Module):
     def loss(self, mask_pred, mask_targets, labels):
         loss = dict()
         if mask_pred.size(0) == 0:
-            loss_mask = mask_pred.sum() * 0
+            loss_mask = mask_pred.sum()
         else:
             if self.class_agnostic:
                 loss_mask = self.loss_mask(mask_pred, mask_targets,

mmdet/models/roi_heads/pisa_roi_head.py

@@ -79,9 +79,7 @@ class PISARoIHead(StandardRoIHead):
             mask_results = self._mask_forward_train(x, sampling_results,
                                                     bbox_results['bbox_feats'],
                                                     gt_masks, img_metas)
-            # TODO: Support empty tensor input. #2280
-            if mask_results['loss_mask'] is not None:
-                losses.update(mask_results['loss_mask'])
+            losses.update(mask_results['loss_mask'])
 
         return losses
 

mmdet/models/roi_heads/standard_roi_head.py

@@ -126,9 +126,7 @@ class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
             mask_results = self._mask_forward_train(x, sampling_results,
                                                     bbox_results['bbox_feats'],
                                                     gt_masks, img_metas)
-            # TODO: Support empty tensor input. #2280
-            if mask_results['loss_mask'] is not None:
-                losses.update(mask_results['loss_mask'])
+            losses.update(mask_results['loss_mask'])
 
         return losses
 
@@ -166,8 +164,6 @@ class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
         training."""
         if not self.share_roi_extractor:
             pos_rois = bbox2roi([res.pos_bboxes for res in sampling_results])
-            if pos_rois.shape[0] == 0:
-                return dict(loss_mask=None)
             mask_results = self._mask_forward(x, pos_rois)
         else:
             pos_inds = []
@@ -184,8 +180,7 @@ class StandardRoIHead(BaseRoIHead, BBoxTestMixin, MaskTestMixin):
                         device=device,
                         dtype=torch.uint8))
             pos_inds = torch.cat(pos_inds)
-            if pos_inds.shape[0] == 0:
-                return dict(loss_mask=None)
+
             mask_results = self._mask_forward(
                 x, pos_inds=pos_inds, bbox_feats=bbox_feats)