Repository: open-mmlab/mmdetection

Issue:
Error while train default Faster R-CNN model
RuntimeError: operation does not have an identity.
It can work sometimes
![2021-05-24 15-06-54屏幕截图](https://user-images.githubusercontent.com/53067819/119310122-caf31800-bc1b-11eb-9748-2622cf5be0cf.png)

Error while train default Faster R-CNN model
RuntimeError: operation does not have an identity.
It can work sometimes
![2021-05-24 15-06-54屏幕截图](https://user-images.githubusercontent.com/53067819/119310122-caf31800-bc1b-11eb-9748-2622cf5be0cf.png)



Discussion:

Resolving patch:
diff --git a/mmdet/models/dense_heads/rpn_head.py b/mmdet/models/dense_heads/rpn_head.py
--- a/mmdet/models/dense_heads/rpn_head.py
+++ b/mmdet/models/dense_heads/rpn_head.py
@@ -4,6 +4,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from mmcv.ops import batched_nms
+from mmcv.runner import force_fp32
 
 from ..builder import HEADS
 from .anchor_head import AnchorHead
@@ -74,30 +75,29 @@ def loss(self,
         return dict(
             loss_rpn_cls=losses['loss_cls'], loss_rpn_bbox=losses['loss_bbox'])
 
-    def _get_bboxes(self,
-                    cls_scores,
-                    bbox_preds,
-                    mlvl_anchors,
-                    img_shapes,
-                    scale_factors,
-                    cfg,
-                    rescale=False):
-        """Transform outputs for a single batch item into bbox predictions.
+    @force_fp32(apply_to=('cls_scores', 'bbox_preds'))
+    def get_bboxes(self,
+                   cls_scores,
+                   bbox_preds,
+                   img_metas,
+                   cfg=None,
+                   rescale=False,
+                   with_nms=True):
+        """Transform network output for a batch into bbox predictions.
 
         Args:
             cls_scores (list[Tensor]): Box scores for each scale level
-                Has shape (N, num_anchors * num_classes, H, W).
+                Has shape (N, num_anchors * num_classes, H, W)
             bbox_preds (list[Tensor]): Box energies / deltas for each scale
-                level with shape (N, num_anchors * 4, H, W).
-            mlvl_anchors (list[Tensor]): Box reference for each scale level
-                with shape (num_total_anchors, 4).
-            img_shapes (list[tuple[int]]): Shape of the input image,
-                (height, width, 3).
-            scale_factors (list[ndarray]): Scale factor of the image arange as
-                (w_scale, h_scale, w_scale, h_scale).
-            cfg (mmcv.Config): Test / postprocessing configuration,
-                if None, test_cfg would be used.
+                level with shape (N, num_anchors * 4, H, W)
+            img_metas (list[dict]): Meta information of each image, e.g.,
+                image size, scaling factor, etc.
+            cfg (mmcv.Config | None): Test / postprocessing configuration,
+                if None, test_cfg would be used
             rescale (bool): If True, return boxes in original image space.
+                Default: False.
+            with_nms (bool): If True, do nms before return boxes.
+                Default: True.
 
         Returns:
             list[tuple[Tensor, Tensor]]: Each item in result_list is 2-tuple.
@@ -107,6 +107,61 @@ def _get_bboxes(self,
                 (n,) tensor where each item is the predicted class label of the
                 corresponding box.
         """
+        assert with_nms, '``with_nms`` in RPNHead should always True'
+        assert len(cls_scores) == len(bbox_preds)
+        num_levels = len(cls_scores)
+        device = cls_scores[0].device
+        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
+        mlvl_anchors = self.anchor_generator.grid_anchors(
+            featmap_sizes, device=device)
+
+        result_list = []
+        for img_id in range(len(img_metas)):
+            cls_score_list = [
+                cls_scores[i][img_id].detach() for i in range(num_levels)
+            ]
+            bbox_pred_list = [
+                bbox_preds[i][img_id].detach() for i in range(num_levels)
+            ]
+            img_shape = img_metas[img_id]['img_shape']
+            scale_factor = img_metas[img_id]['scale_factor']
+            proposals = self._get_bboxes_single(cls_score_list, bbox_pred_list,
+                                                mlvl_anchors, img_shape,
+                                                scale_factor, cfg, rescale)
+            result_list.append(proposals)
+        return result_list
+
+    def _get_bboxes_single(self,
+                           cls_scores,
+                           bbox_preds,
+                           mlvl_anchors,
+                           img_shape,
+                           scale_factor,
+                           cfg,
+                           rescale=False):
+        """Transform outputs for a single batch item into bbox predictions.
+
+          Args:
+            cls_scores (list[Tensor]): Box scores of all scale level
+                each item has shape (num_anchors * num_classes, H, W).
+            bbox_preds (list[Tensor]): Box energies / deltas of all
+                scale level, each item has shape (num_anchors * 4, H, W).
+            mlvl_anchors (list[Tensor]): Anchors of all scale level
+                each item has shape (num_total_anchors, 4).
+            img_shape (tuple[int]): Shape of the input image,
+                (height, width, 3).
+            scale_factor (ndarray): Scale factor of the image arrange as
+                (w_scale, h_scale, w_scale, h_scale).
+            cfg (mmcv.Config): Test / postprocessing configuration,
+                if None, test_cfg would be used.
+            rescale (bool): If True, return boxes in original image space.
+                Default: False.
+
+        Returns:
+            Tensor: Labeled boxes in shape (n, 5), where the first 4 columns
+                are bounding box positions (tl_x, tl_y, br_x, br_y) and the
+                5-th column is a score between 0 and 1.
+        """
         cfg = self.test_cfg if cfg is None else cfg
         cfg = copy.deepcopy(cfg)
         # bboxes from different level should be independent during NMS,
@@ -115,77 +170,59 @@ def _get_bboxes(self,
         mlvl_scores = []
         mlvl_bbox_preds = []
         mlvl_valid_anchors = []
-        batch_size = cls_scores[0].shape[0]
         for idx in range(len(cls_scores)):
             rpn_cls_score = cls_scores[idx]
             rpn_bbox_pred = bbox_preds[idx]
             assert rpn_cls_score.size()[-2:] == rpn_bbox_pred.size()[-2:]
-            rpn_cls_score = rpn_cls_score.permute(0, 2, 3, 1)
+            rpn_cls_score = rpn_cls_score.permute(1, 2, 0)
             if self.use_sigmoid_cls:
-                rpn_cls_score = rpn_cls_score.reshape(batch_size, -1)
+                rpn_cls_score = rpn_cls_score.reshape(-1)
                 scores = rpn_cls_score.sigmoid()
             else:
-                rpn_cls_score = rpn_cls_score.reshape(batch_size, -1, 2)
+                rpn_cls_score = rpn_cls_score.reshape(-1, 2)
                 # We set FG labels to [0, num_class-1] and BG label to
                 # num_class in RPN head since mmdet v2.5, which is unified to
                 # be consistent with other head since mmdet v2.0. In mmdet v2.0
                 # to v2.4 we keep BG label as 0 and FG label as 1 in rpn head.
-                scores = rpn_cls_score.softmax(-1)[..., 0]
-            rpn_bbox_pred = rpn_bbox_pred.permute(0, 2, 3, 1).reshape(
-                batch_size, -1, 4)
+                scores = rpn_cls_score.softmax(dim=1)[:, 0]
+            rpn_bbox_pred = rpn_bbox_pred.permute(1, 2, 0).reshape(-1, 4)
             anchors = mlvl_anchors[idx]
-            anchors = anchors.expand_as(rpn_bbox_pred)
-            nms_pre = cfg.get('nms_pre', -1)
-            if nms_pre > 0 and rpn_bbox_pred.size(1) > nms_pre:
+            if cfg.nms_pre > 0 and scores.shape[0] > cfg.nms_pre:
                 # sort is faster than topk
+                # _, topk_inds = scores.topk(cfg.nms_pre)
                 ranked_scores, rank_inds = scores.sort(descending=True)
-                topk_inds = rank_inds[:, :cfg.nms_pre]
-                scores = ranked_scores[:, :cfg.nms_pre]
-                batch_inds = torch.arange(batch_size).view(
-                    -1, 1).expand_as(topk_inds)
-                rpn_bbox_pred = rpn_bbox_pred[batch_inds, topk_inds, :]
-                anchors = anchors[batch_inds, topk_inds, :]
-
+                topk_inds = rank_inds[:cfg.nms_pre]
+                scores = ranked_scores[:cfg.nms_pre]
+                rpn_bbox_pred = rpn_bbox_pred[topk_inds, :]
+                anchors = anchors[topk_inds, :]
             mlvl_scores.append(scores)
             mlvl_bbox_preds.append(rpn_bbox_pred)
             mlvl_valid_anchors.append(anchors)
             level_ids.append(
-                scores.new_full((
-                    batch_size,
-                    scores.size(1),
-                ),
-                                idx,
-                                dtype=torch.long))
+                scores.new_full((scores.size(0), ), idx, dtype=torch.long))
 
-        batch_mlvl_scores = torch.cat(mlvl_scores, dim=1)
-        batch_mlvl_anchors = torch.cat(mlvl_valid_anchors, dim=1)
-        batch_mlvl_rpn_bbox_pred = torch.cat(mlvl_bbox_preds, dim=1)
-        batch_mlvl_proposals = self.bbox_coder.decode(
-            batch_mlvl_anchors, batch_mlvl_rpn_bbox_pred, max_shape=img_shapes)
-        batch_mlvl_ids = torch.cat(level_ids, dim=1)
+        scores = torch.cat(mlvl_scores)
+        anchors = torch.cat(mlvl_valid_anchors)
+        rpn_bbox_pred = torch.cat(mlvl_bbox_preds)
+        proposals = self.bbox_coder.decode(
+            anchors, rpn_bbox_pred, max_shape=img_shape)
+        ids = torch.cat(level_ids)
 
-        result_list = []
-        for (mlvl_proposals, mlvl_scores,
-             mlvl_ids) in zip(batch_mlvl_proposals, batch_mlvl_scores,
-                              batch_mlvl_ids):
-            if cfg.min_bbox_size >= 0:
-                w = mlvl_proposals[:, 2] - mlvl_proposals[:, 0]
-                h = mlvl_proposals[:, 3] - mlvl_proposals[:, 1]
-                valid_ind = torch.nonzero(
-                    (w > cfg.min_bbox_size)
-                    & (h > cfg.min_bbox_size),
-                    as_tuple=False).squeeze()
-                if valid_ind.sum().item() != len(mlvl_proposals):
-                    mlvl_proposals = mlvl_proposals[valid_ind, :]
-                    mlvl_scores = mlvl_scores[valid_ind]
-                    mlvl_ids = mlvl_ids[valid_ind]
-
-            dets, keep = batched_nms(mlvl_proposals, mlvl_scores, mlvl_ids,
-                                     cfg.nms)
-            result_list.append(dets[:cfg.max_per_img])
-        return result_list
+        if cfg.min_bbox_size > 0:
+            w = proposals[:, 2] - proposals[:, 0]
+            h = proposals[:, 3] - proposals[:, 1]
+            valid_mask = (w >= cfg.min_bbox_size) & (h >= cfg.min_bbox_size)
+            if not valid_mask.all():
+                proposals = proposals[valid_mask]
+                scores = scores[valid_mask]
+                ids = ids[valid_mask]
+        if proposals.numel() > 0:
+            dets, keep = batched_nms(proposals, scores, ids, cfg.nms)
+        else:
+            return proposals.new_zeros(0, 5)
+
+        return dets[:cfg.max_per_img]
 
-    # TODO: waiting for refactor the anchor_head and anchor_free head
     def onnx_export(self, x, img_metas):
         """Test without augmentation.
 
diff --git a/mmdet/models/roi_heads/bbox_heads/bbox_head.py b/mmdet/models/roi_heads/bbox_heads/bbox_head.py
--- a/mmdet/models/roi_heads/bbox_heads/bbox_head.py
+++ b/mmdet/models/roi_heads/bbox_heads/bbox_head.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -317,122 +316,59 @@ def get_bboxes(self,
                    cfg=None):
         """Transform network output for a batch into bbox predictions.
 
-        In most case except Cascade R-CNN, HTC, AugTest..,
-        the dimensions of input rois, cls_score, bbox_pred are equal
-        to 3, and batch dimension is the first dimension, for example
-        roi has shape (B, num_boxes, 5), return is a
-        tuple[list[Tensor], list[Tensor]],
-        the length of list in tuple is equal to the batch_size.
-        otherwise, the input tensor has only 2 dimensions,
-        and return is a tuple[Tensor, Tensor].
-
         Args:
-            rois (Tensor): Boxes to be transformed. Has shape (num_boxes, 5)
-               or (B, num_boxes, 5)
-            cls_score (Tensor): Box scores, Has shape
-               (B, num_boxes, num_classes + 1) in `batch_model`, otherwise
-                has shape (num_boxes, num_classes + 1).
-            bbox_pred (Tensor, optional): Box energies / deltas. Has shape
-                (B, num_boxes, num_classes * 4) in `batch_model`, otherwise
+            rois (Tensor): Boxes to be transformed. Has shape (num_boxes, 5).
+                last dimension 5 arrange as (batch_index, x1, y1, x2, y2).
+            cls_score (Tensor): Box scores, has shape
+                (num_boxes, num_classes + 1).
+            bbox_pred (Tensor, optional): Box energies / deltas.
                 has shape (num_boxes, num_classes * 4).
-            img_shape (Sequence[int] or Sequence[
-                Sequence[int]], optional): Maximum bounds for boxes, specifies
-                (H, W, C) or (H, W). If rois shape is (B, num_boxes, 4), then
-                the max_shape should be a Sequence[Sequence[int]]
-                and the length of max_shape should be equal to the batch_size.
-            scale_factor (tuple[ndarray] or ndarray): Scale factor of the
-               image arrange as (w_scale, h_scale, w_scale, h_scale). In
-               `batch_mode`, the scale_factor shape is tuple[ndarray].
-               the length should be equal to the batch size.
+            img_shape (Sequence[int], optional): Maximum bounds for boxes,
+                specifies (H, W, C) or (H, W).
+            scale_factor (ndarray): Scale factor of the
+               image arrange as (w_scale, h_scale, w_scale, h_scale).
             rescale (bool): If True, return boxes in original image space.
                 Default: False.
             cfg (obj:`ConfigDict`): `test_cfg` of Bbox Head. Default: None
 
         Returns:
-            tuple[list[Tensor], list[Tensor]] or tuple[Tensor, Tensor]:
-                If the input has a batch dimension, the return value is
-                a tuple of the list. The first list contains the boxes of
-                the corresponding image in a batch, each tensor has the
-                shape (num_boxes, 5) and last dimension 5 represent
-                (tl_x, tl_y, br_x, br_y, score). Each Tensor in the second
-                list is the labels with shape (num_boxes, ). The length of
-                both lists should be equal to batch_size. Otherwise return
-                value is a tuple of two tensors, the first tensor is the
-                boxes with scores, the second tensor is the labels, both
-                have the same shape as the first case.
+            tuple[Tensor, Tensor]:
+                Fisrt tensor is `det_bboxes`, has the shape
+                (num_boxes, 5) and last
+                dimension 5 represent (tl_x, tl_y, br_x, br_y, score).
+                Second tensor is the labels with shape (num_boxes, ).
         """
 
-        # TODO: revert to single image inference
-
         # some loss (Seesaw loss..) may have custom activation
         if self.custom_cls_channels:
             scores = self.loss_cls.get_activation(cls_score)
         else:
             scores = F.softmax(
                 cls_score, dim=-1) if cls_score is not None else None
-
-        if rois.ndim == 2:
-            # e.g. AugTest, Cascade R-CNN, HTC, SCNet...
-            batch_mode = False
-            # add batch dimension
-            if scores is not None:
-                scores = scores.unsqueeze(0)
-            if bbox_pred is not None:
-                bbox_pred = bbox_pred.unsqueeze(0)
-            rois = rois.unsqueeze(0)
-
-            assert isinstance(scale_factor, np.ndarray)
-            scale_factor = (scale_factor, )
-
-        elif rois.ndim == 3:
-            # all input tensor have batch dimension
-            batch_mode = True
-            assert isinstance(scale_factor, tuple)
-        else:
-            raise NotImplementedError(f'Unexpect shape of roi {rois.shape}')
-
         # bbox_pred would be None in some detector when with_reg is False,
         # e.g. Grid R-CNN.
         if bbox_pred is not None:
             bboxes = self.bbox_coder.decode(
                 rois[..., 1:], bbox_pred, max_shape=img_shape)
         else:
-            bboxes = rois[..., 1:].clone()
+            bboxes = rois[:, 1:].clone()
             if img_shape is not None:
-                max_shape = bboxes.new_tensor(img_shape)[..., :2]
-                min_xy = bboxes.new_tensor(0)
-                max_xy = torch.cat(
-                    [max_shape] * 2, dim=-1).flip(-1).unsqueeze(-2)
-                bboxes = torch.where(bboxes < min_xy, min_xy, bboxes)
-                bboxes = torch.where(bboxes > max_xy, max_xy, bboxes)
+                bboxes[:, [0, 2]].clamp_(min=0, max=img_shape[1])
+                bboxes[:, [1, 3]].clamp_(min=0, max=img_shape[0])
 
-        num_bboxes = bboxes.size(-2)
-        if rescale and num_bboxes > 0:
-            # B, 1, bboxes.size(-1)
-            scale_factor = bboxes.new_tensor(scale_factor).unsqueeze(1).repeat(
-                1, 1,
-                bboxes.size(-1) // 4)
-            bboxes /= scale_factor
-
-        det_bboxes = []
-        det_labels = []
-        for (bbox, score) in zip(bboxes, scores):
-            if cfg is not None:
-                det_bbox, det_label = multiclass_nms(bbox, score,
-                                                     cfg.score_thr, cfg.nms,
-                                                     cfg.max_per_img)
-            else:
-                det_bbox, det_label = bbox, score
-            det_bboxes.append(det_bbox)
-            det_labels.append(det_label)
-
-        if not batch_mode:
-            single_det_bboxes = det_bboxes[0]
-            single_det_labels = det_labels[0]
-            # tuple[Tensor, Tensor]
-            return single_det_bboxes, single_det_labels
+        if rescale and bboxes.size(0) > 0:
+
+            scale_factor = bboxes.new_tensor(scale_factor)
+            bboxes = (bboxes.view(bboxes.size(0), -1, 4) / scale_factor).view(
+                bboxes.size()[0], -1)
+
+        if cfg is None:
+            return bboxes, scores
         else:
-            # tuple[list[Tensor], list[Tensor]]
+            det_bboxes, det_labels = multiclass_nms(bboxes, scores,
+                                                    cfg.score_thr, cfg.nms,
+                                                    cfg.max_per_img)
+
             return det_bboxes, det_labels
 
     @force_fp32(apply_to=('bbox_preds', ))
diff --git a/mmdet/models/roi_heads/cascade_roi_head.py b/mmdet/models/roi_heads/cascade_roi_head.py
--- a/mmdet/models/roi_heads/cascade_roi_head.py
+++ b/mmdet/models/roi_heads/cascade_roi_head.py
@@ -1,3 +1,4 @@
+import numpy as np
 import torch
 import torch.nn as nn
 from mmcv.runner import ModuleList
@@ -481,13 +482,14 @@ def aug_test(self, features, proposal_list, img_metas, rescale=False):
                                                self.test_cfg)
 
                 ori_shape = img_metas[0][0]['ori_shape']
+                dummy_scale_factor = np.ones(4)
                 segm_result = self.mask_head[-1].get_seg_masks(
                     merged_masks,
                     det_bboxes,
                     det_labels,
                     rcnn_test_cfg,
                     ori_shape,
-                    scale_factor=1.0,
+                    scale_factor=dummy_scale_factor,
                     rescale=False)
             return [(bbox_result, segm_result)]
         else:


Test patch:
diff --git a/mmdet/models/roi_heads/test_mixins.py b/mmdet/models/roi_heads/test_mixins.py
--- a/mmdet/models/roi_heads/test_mixins.py
+++ b/mmdet/models/roi_heads/test_mixins.py
@@ -1,6 +1,7 @@
 import logging
 import sys
 
+import numpy as np
 import torch
 
 from mmdet.core import (bbox2roi, bbox_mapping, merge_aug_bboxes,
@@ -71,88 +72,46 @@ def simple_test_bboxes(self,
                 in the second list is the labels with shape (num_boxes, ).
                 The length of both lists should be equal to batch_size.
         """
-        # get origin input shape to support onnx dynamic input shape
 
+        rois = bbox2roi(proposals)
+        bbox_results = self._bbox_forward(x, rois)
         img_shapes = tuple(meta['img_shape'] for meta in img_metas)
         scale_factors = tuple(meta['scale_factor'] for meta in img_metas)
 
-        # The length of proposals of different batches may be different.
-        # In order to form a batch, a padding operation is required.
-        max_size = max([proposal.size(0) for proposal in proposals])
-        # padding to form a batch
-        for i, proposal in enumerate(proposals):
-            supplement = proposal.new_full(
-                (max_size - proposal.size(0), proposal.size(1)), 0)
-            proposals[i] = torch.cat((supplement, proposal), dim=0)
-        rois = torch.stack(proposals, dim=0)
-
-        batch_index = torch.arange(
-            rois.size(0), device=rois.device).float().view(-1, 1, 1).expand(
-                rois.size(0), rois.size(1), 1)
-        rois = torch.cat([batch_index, rois[..., :4]], dim=-1)
-        batch_size = rois.shape[0]
-        num_proposals_per_img = rois.shape[1]
-
-        # Eliminate the batch dimension
-        rois = rois.view(-1, 5)
-        bbox_results = self._bbox_forward(x, rois)
+        # split batch bbox prediction back to each image
         cls_score = bbox_results['cls_score']
         bbox_pred = bbox_results['bbox_pred']
+        num_proposals_per_img = tuple(len(p) for p in proposals)
+        rois = rois.split(num_proposals_per_img, 0)
+        cls_score = cls_score.split(num_proposals_per_img, 0)
 
-        # Recover the batch dimension
-        rois = rois.reshape(batch_size, num_proposals_per_img, rois.size(-1))
-        cls_score = cls_score.reshape(batch_size, num_proposals_per_img,
-                                      cls_score.size(-1))
-
-        # remove padding, ignore batch_index when calculating mask
-        supplement_mask = rois.abs()[..., 1:].sum(dim=-1) == 0
-        cls_score[supplement_mask, :] = 0
-
-        # bbox_pred would be None in some detector when with_reg is False,
-        # e.g. Grid R-CNN.
+        # some detector with_reg is False, bbox_pred will be None
         if bbox_pred is not None:
+            # TODO move this to a sabl_roi_head
             # the bbox prediction of some detectors like SABL is not Tensor
             if isinstance(bbox_pred, torch.Tensor):
-                bbox_pred = bbox_pred.reshape(batch_size,
-                                              num_proposals_per_img,
-                                              bbox_pred.size(-1))
-                bbox_pred[supplement_mask, :] = 0
+                bbox_pred = bbox_pred.split(num_proposals_per_img, 0)
             else:
-                # TODO: Looking forward to a better way
-                # TODO move these special process to a corresponding head
-                # For SABL
-                bbox_preds = self.bbox_head.bbox_pred_split(
+                bbox_pred = self.bbox_head.bbox_pred_split(
                     bbox_pred, num_proposals_per_img)
-                # apply bbox post-processing to each image individually
-                det_bboxes = []
-                det_labels = []
-                for i in range(len(proposals)):
-                    # remove padding
-                    supplement_mask = proposals[i].abs().sum(dim=-1) == 0
-                    for bbox in bbox_preds[i]:
-                        bbox[supplement_mask] = 0
-                    det_bbox, det_label = self.bbox_head.get_bboxes(
-                        rois[i],
-                        cls_score[i],
-                        bbox_preds[i],
-                        img_shapes[i],
-                        scale_factors[i],
-                        rescale=rescale,
-                        cfg=rcnn_test_cfg)
-                    det_bboxes.append(det_bbox)
-                    det_labels.append(det_label)
-                return det_bboxes, det_labels
         else:
-            bbox_pred = None
-
-        return self.bbox_head.get_bboxes(
-            rois,
-            cls_score,
-            bbox_pred,
-            img_shapes,
-            scale_factors,
-            rescale=rescale,
-            cfg=rcnn_test_cfg)
+            bbox_pred = (None, ) * len(proposals)
+
+        # apply bbox post-processing to each image individually
+        det_bboxes = []
+        det_labels = []
+        for i in range(len(proposals)):
+            det_bbox, det_label = self.bbox_head.get_bboxes(
+                rois[i],
+                cls_score[i],
+                bbox_pred[i],
+                img_shapes[i],
+                scale_factors[i],
+                rescale=rescale,
+                cfg=rcnn_test_cfg)
+            det_bboxes.append(det_bbox)
+            det_labels.append(det_label)
+        return det_bboxes, det_labels
 
     def aug_test_bboxes(self, feats, img_metas, proposal_list, rcnn_test_cfg):
         """Test det bboxes with test time augmentation."""
@@ -207,7 +166,8 @@ async def async_test_mask(self,
             if det_bboxes.shape[0] == 0:
                 segm_result = [[] for _ in range(self.mask_head.num_classes)]
             else:
-                if rescale:
+                if rescale and not isinstance(scale_factor,
+                                              (float, torch.Tensor)):
                     scale_factor = det_bboxes.new_tensor(scale_factor)
                 _bboxes = (
                     det_bboxes[:, :4] *
@@ -244,69 +204,50 @@ def simple_test_mask(self,
         ori_shapes = tuple(meta['ori_shape'] for meta in img_metas)
         scale_factors = tuple(meta['scale_factor'] for meta in img_metas)
 
+        if isinstance(scale_factors[0], float):
+            logger.warning(
+                'Scale factor in img_metas should be a '
+                'ndarray with shape (4,) '
+                'arrange as (factor_w, factor_h, factor_w, factor_h), '
+                'The scale_factor with float type has been deprecated. ')
+            scale_factors = np.array([scale_factors] * 4)
+
+        num_imgs = len(det_bboxes)
         if all(det_bbox.shape[0] == 0 for det_bbox in det_bboxes):
             segm_results = [[[] for _ in range(self.mask_head.num_classes)]
-                            for _ in range(len(det_bboxes))]
-            return segm_results
-
-        # The length of proposals of different batches may be different.
-        # In order to form a batch, a padding operation is required.
-
-        # padding to form a batch
-        max_size = max([bboxes.size(0) for bboxes in det_bboxes])
-        for i, (bbox, label) in enumerate(zip(det_bboxes, det_labels)):
-            supplement_bbox = bbox.new_full(
-                (max_size - bbox.size(0), bbox.size(1)), 0)
-            supplement_label = label.new_full((max_size - label.size(0), ), 0)
-            det_bboxes[i] = torch.cat((supplement_bbox, bbox), dim=0)
-            det_labels[i] = torch.cat((supplement_label, label), dim=0)
-        det_bboxes = torch.stack(det_bboxes, dim=0)
-        det_labels = torch.stack(det_labels, dim=0)
-
-        batch_size = det_bboxes.size(0)
-        num_proposals_per_img = det_bboxes.shape[1]
-
-        # if det_bboxes is rescaled to the original image size, we need to
-        # rescale it back to the testing scale to obtain RoIs.
-        det_bboxes = det_bboxes[..., :4]
-        if rescale:
-            scale_factors = det_bboxes.new_tensor(scale_factors)
-            det_bboxes = det_bboxes * scale_factors.unsqueeze(1)
-
-        batch_index = torch.arange(
-            det_bboxes.size(0), device=det_bboxes.device).float().view(
-                -1, 1, 1).expand(det_bboxes.size(0), det_bboxes.size(1), 1)
-        mask_rois = torch.cat([batch_index, det_bboxes], dim=-1)
-        mask_rois = mask_rois.view(-1, 5)
-        mask_results = self._mask_forward(x, mask_rois)
-        mask_pred = mask_results['mask_pred']
-
-        # Recover the batch dimension
-        mask_preds = mask_pred.reshape(batch_size, num_proposals_per_img,
-                                       *mask_pred.shape[1:])
-
-        # apply mask post-processing to each image individually
-        segm_results = []
-        for i in range(batch_size):
-            mask_pred = mask_preds[i]
-            det_bbox = det_bboxes[i]
-            det_label = det_labels[i]
-
-            # remove padding
-            supplement_mask = det_bbox.abs().sum(dim=-1) != 0
-            mask_pred = mask_pred[supplement_mask]
-            det_bbox = det_bbox[supplement_mask]
-            det_label = det_label[supplement_mask]
-
-            if det_label.shape[0] == 0:
-                segm_results.append([[]
-                                     for _ in range(self.mask_head.num_classes)
-                                     ])
-            else:
-                segm_result = self.mask_head.get_seg_masks(
-                    mask_pred, det_bbox, det_label, self.test_cfg,
-                    ori_shapes[i], scale_factors[i], rescale)
-                segm_results.append(segm_result)
+                            for _ in range(num_imgs)]
+        else:
+            # if det_bboxes is rescaled to the original image size, we need to
+            # rescale it back to the testing scale to obtain RoIs.
+            if rescale:
+                scale_factors = [
+                    torch.from_numpy(scale_factor).to(det_bboxes[0].device)
+                    for scale_factor in scale_factors
+                ]
+            _bboxes = [
+                det_bboxes[i][:, :4] *
+                scale_factors[i] if rescale else det_bboxes[i][:, :4]
+                for i in range(len(det_bboxes))
+            ]
+            mask_rois = bbox2roi(_bboxes)
+            mask_results = self._mask_forward(x, mask_rois)
+            mask_pred = mask_results['mask_pred']
+            # split batch mask prediction back to each image
+            num_mask_roi_per_img = [len(det_bbox) for det_bbox in det_bboxes]
+            mask_preds = mask_pred.split(num_mask_roi_per_img, 0)
+
+            # apply mask post-processing to each image individually
+            segm_results = []
+            for i in range(num_imgs):
+                if det_bboxes[i].shape[0] == 0:
+                    segm_results.append(
+                        [[] for _ in range(self.mask_head.num_classes)])
+                else:
+                    segm_result = self.mask_head.get_seg_masks(
+                        mask_preds[i], _bboxes[i], det_labels[i],
+                        self.test_cfg, ori_shapes[i], scale_factors[i],
+                        rescale)
+                    segm_results.append(segm_result)
         return segm_results
 
     def aug_test_mask(self, feats, img_metas, det_bboxes, det_labels):
diff --git a/tests/test_models/test_roi_heads/test_bbox_head.py b/tests/test_models/test_roi_heads/test_bbox_head.py
--- a/tests/test_models/test_roi_heads/test_bbox_head.py
+++ b/tests/test_models/test_roi_heads/test_bbox_head.py
@@ -66,8 +66,8 @@ def test_bbox_head_loss():
     assert losses.get('loss_bbox', 0) > 0, 'box-loss should be non-zero'
 
 
-@pytest.mark.parametrize(['num_sample', 'num_batch'], [[2, 2], [0, 2], [0, 0]])
-def test_bbox_head_get_bboxes(num_sample, num_batch):
+@pytest.mark.parametrize('num_sample', [0, 1, 2])
+def test_bbox_head_get_bboxes(num_sample):
     self = BBoxHead(reg_class_agnostic=True)
 
     num_class = 6
@@ -84,14 +84,6 @@ def test_bbox_head_get_bboxes(num_sample, num_batch):
         assert det_bboxes.shape == bbox_pred.shape
         assert det_labels.shape == cls_score.shape
 
-    rois = torch.rand((num_batch, num_sample, 5))
-    cls_score = torch.rand((num_batch, num_sample, num_class))
-    bbox_pred = torch.rand((num_batch, num_sample, 4))
-    scale_factor = (np.array([2.0, 2.0, 2.0, 2.0]), ) * num_batch
-    det_bboxes, det_labels = self.get_bboxes(
-        rois, cls_score, bbox_pred, None, scale_factor, rescale=True)
-    assert len(det_bboxes) == num_batch and len(det_labels) == num_batch
-
 
 def test_refine_boxes():
     """Mirrors the doctest in