Support empty tensor input for some models. (open-mmlab#2280)

* support-empty-tensor

* minor update

* add unit test

* fix unit test

* add assert value equal

* simplify some codes

* simplify unit tests

* distinguish x with x_empty and x_normal

* ref only forward once

* fix python3.5 ci error

Author: yhcao6

mmdet/core/mask/mask_target.py

@@ -8,7 +8,9 @@ def mask_target(pos_proposals_list, pos_assigned_gt_inds_list, gt_masks_list,
     cfg_list = [cfg for _ in range(len(pos_proposals_list))]
     mask_targets = map(mask_target_single, pos_proposals_list,
                        pos_assigned_gt_inds_list, gt_masks_list, cfg_list)
-    mask_targets = torch.cat(list(mask_targets))
+    mask_targets = list(mask_targets)
+    if len(mask_targets) > 0:
+        mask_targets = torch.cat(mask_targets)
     return mask_targets
 
 

mmdet/models/bbox_heads/bbox_head.py

@@ -138,6 +138,8 @@ def loss(self,
                     bbox_weights[pos_inds.type(torch.bool)],
                     avg_factor=bbox_targets.size(0),
                     reduction_override=reduction_override)
+            else:
+                losses['loss_bbox'] = bbox_pred.sum() * 0
         return losses
 
     @force_fp32(apply_to=('cls_score', 'bbox_pred'))

mmdet/models/mask_heads/fcn_mask_head.py

@@ -5,7 +5,7 @@
 from torch.nn.modules.utils import _pair
 
 from mmdet.core import auto_fp16, force_fp32, mask_target
-from mmdet.ops import ConvModule, build_upsample_layer
+from mmdet.ops import Conv2d, ConvModule, build_upsample_layer
 from mmdet.ops.carafe import CARAFEPack
 from mmdet.ops.grid_sampler import grid_sample
 from ..builder import build_loss
@@ -98,7 +98,7 @@ def __init__(self,
         logits_in_channel = (
             self.conv_out_channels
             if self.upsample_method == 'deconv' else upsample_in_channels)
-        self.conv_logits = nn.Conv2d(logits_in_channel, out_channels, 1)
+        self.conv_logits = Conv2d(logits_in_channel, out_channels, 1)
         self.relu = nn.ReLU(inplace=True)
         self.debug_imgs = None
 
@@ -136,11 +136,14 @@ def get_target(self, sampling_results, gt_masks, rcnn_train_cfg):
     @force_fp32(apply_to=('mask_pred', ))
     def loss(self, mask_pred, mask_targets, labels):
         loss = dict()
-        if self.class_agnostic:
-            loss_mask = self.loss_mask(mask_pred, mask_targets,
-                                       torch.zeros_like(labels))
+        if mask_pred.size(0) == 0:
+            loss_mask = mask_pred.sum() * 0
         else:
-            loss_mask = self.loss_mask(mask_pred, mask_targets, labels)
+            if self.class_agnostic:
+                loss_mask = self.loss_mask(mask_pred, mask_targets,
+                                           torch.zeros_like(labels))
+            else:
+                loss_mask = self.loss_mask(mask_pred, mask_targets, labels)
         loss['loss_mask'] = loss_mask
         return loss
 

mmdet/models/mask_heads/maskiou_head.py

@@ -5,6 +5,7 @@
 from torch.nn.modules.utils import _pair
 
 from mmdet.core import force_fp32
+from mmdet.ops import Conv2d, Linear, MaxPool2d
 from ..builder import build_loss
 from ..registry import HEADS
 
@@ -41,7 +42,7 @@ def __init__(self,
                 in_channels = self.conv_out_channels
             stride = 2 if i == num_convs - 1 else 1
             self.convs.append(
-                nn.Conv2d(
+                Conv2d(
                     in_channels,
                     self.conv_out_channels,
                     3,
@@ -55,11 +56,11 @@ def __init__(self,
             in_channels = (
                 self.conv_out_channels *
                 pooled_area if i == 0 else self.fc_out_channels)
-            self.fcs.append(nn.Linear(in_channels, self.fc_out_channels))
+            self.fcs.append(Linear(in_channels, self.fc_out_channels))
 
-        self.fc_mask_iou = nn.Linear(self.fc_out_channels, self.num_classes)
+        self.fc_mask_iou = Linear(self.fc_out_channels, self.num_classes)
         self.relu = nn.ReLU()
-        self.max_pool = nn.MaxPool2d(2, 2)
+        self.max_pool = MaxPool2d(2, 2)
         self.loss_iou = build_loss(loss_iou)
 
     def init_weights(self):
@@ -82,7 +83,7 @@ def forward(self, mask_feat, mask_pred):
 
         for conv in self.convs:
             x = self.relu(conv(x))
-        x = x.view(x.size(0), -1)
+        x = x.flatten(1)
         for fc in self.fcs:
             x = self.relu(fc(x))
         mask_iou = self.fc_mask_iou(x)
@@ -95,7 +96,7 @@ def loss(self, mask_iou_pred, mask_iou_targets):
             loss_mask_iou = self.loss_iou(mask_iou_pred[pos_inds],
                                           mask_iou_targets[pos_inds])
         else:
-            loss_mask_iou = mask_iou_pred * 0
+            loss_mask_iou = mask_iou_pred.sum() * 0
         return dict(loss_mask_iou=loss_mask_iou)
 
     @force_fp32(apply_to=('mask_pred', ))

mmdet/models/roi_extractors/single_level.py

@@ -88,14 +88,17 @@ def roi_rescale(self, rois, scale_factor):
 
     @force_fp32(apply_to=('feats', ), out_fp16=True)
     def forward(self, feats, rois, roi_scale_factor=None):
-        if len(feats) == 1:
-            return self.roi_layers[0](feats[0], rois)
-
         out_size = self.roi_layers[0].out_size
         num_levels = len(feats)
-        target_lvls = self.map_roi_levels(rois, num_levels)
         roi_feats = feats[0].new_zeros(
             rois.size(0), self.out_channels, *out_size)
+
+        if num_levels == 1:
+            if len(rois) == 0:
+                return roi_feats
+            return self.roi_layers[0](feats[0], rois)
+
+        target_lvls = self.map_roi_levels(rois, num_levels)
         if roi_scale_factor is not None:
             rois = self.roi_rescale(rois, roi_scale_factor)
         for i in range(num_levels):

mmdet/ops/__init__.py

@@ -18,6 +18,7 @@
 from .sigmoid_focal_loss import SigmoidFocalLoss, sigmoid_focal_loss
 from .upsample import build_upsample_layer
 from .utils import get_compiler_version, get_compiling_cuda_version
+from .wrappers import Conv2d, ConvTranspose2d, Linear, MaxPool2d
 
 __all__ = [
     'nms', 'soft_nms', 'RoIAlign', 'roi_align', 'RoIPool', 'roi_pool',
@@ -28,5 +29,6 @@
     'MaskedConv2d', 'ContextBlock', 'GeneralizedAttention', 'NonLocal2D',
     'get_compiler_version', 'get_compiling_cuda_version', 'build_conv_layer',
     'ConvModule', 'ConvWS2d', 'conv_ws_2d', 'build_norm_layer', 'Scale',
-    'build_upsample_layer', 'build_plugin_layer', 'batched_nms'
+    'build_upsample_layer', 'build_plugin_layer', 'batched_nms', 'Conv2d',
+    'ConvTranspose2d', 'MaxPool2d', 'Linear'
 ]

mmdet/ops/conv.py

@@ -1,10 +1,9 @@
-from torch import nn as nn
-
 from .conv_ws import ConvWS2d
 from .dcn import DeformConvPack, ModulatedDeformConvPack
+from .wrappers import Conv2d
 
 conv_cfg = {
-    'Conv': nn.Conv2d,
+    'Conv': Conv2d,
     'ConvWS': ConvWS2d,
     'DCN': DeformConvPack,
     'DCNv2': ModulatedDeformConvPack,

mmdet/ops/upsample.py

@@ -3,6 +3,7 @@
 from mmcv.cnn import xavier_init
 
 from .carafe import CARAFEPack
+from .wrappers import ConvTranspose2d
 
 
 class PixelShufflePack(nn.Module):
@@ -45,7 +46,7 @@ def forward(self, x):
     # layer_abbreviation: module
     'nearest': nn.Upsample,
     'bilinear': nn.Upsample,
-    'deconv': nn.ConvTranspose2d,
+    'deconv': ConvTranspose2d,
     'pixel_shuffle': PixelShufflePack,
     'carafe': CARAFEPack
 }

mmdet/ops/wrappers.py

@@ -0,0 +1,98 @@
+"""
+Modified from https://github.com/facebookresearch/detectron2/blob/master
+/detectron2/layers/wrappers.py
+Wrap some nn modules to support empty tensor input.
+Currently, these wrappers are mainly used in mask heads like fcn_mask_head
+and maskiou_heads since mask heads are trained on only positive RoIs.
+"""
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.utils import _pair
+
+
+class NewEmptyTensorOp(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x, new_shape):
+        ctx.shape = x.shape
+        return x.new_empty(new_shape)
+
+    @staticmethod
+    def backward(ctx, grad):
+        shape = ctx.shape
+        return NewEmptyTensorOp.apply(grad, shape), None
+
+
+class Conv2d(nn.Conv2d):
+
+    def forward(self, x):
+        if x.numel() == 0 and torch.__version__ <= '1.4':
+            out_shape = [x.shape[0], self.out_channels]
+            for i, k, p, s, d in zip(x.shape[-2:], self.kernel_size,
+                                     self.padding, self.stride, self.dilation):
+                o = (i + 2 * p - (d * (k - 1) + 1)) // s + 1
+                out_shape.append(o)
+            empty = NewEmptyTensorOp.apply(x, out_shape)
+            if self.training:
+                # produce dummy gradient to avoid DDP warning.
+                dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
+                return empty + dummy
+            else:
+                return empty
+
+        return super().forward(x)
+
+
+class ConvTranspose2d(nn.ConvTranspose2d):
+
+    def forward(self, x):
+        if x.numel() == 0 and torch.__version__ <= '1.4.0':
+            out_shape = [x.shape[0], self.out_channels]
+            for i, k, p, s, d, op in zip(x.shape[-2:], self.kernel_size,
+                                         self.padding, self.stride,
+                                         self.dilation, self.output_padding):
+                out_shape.append((i - 1) * s - 2 * p + (d * (k - 1) + 1) + op)
+            empty = NewEmptyTensorOp.apply(x, out_shape)
+            if self.training:
+                # produce dummy gradient to avoid DDP warning.
+                dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
+                return empty + dummy
+            else:
+                return empty
+
+        return super(ConvTranspose2d, self).forward(x)
+
+
+class MaxPool2d(nn.MaxPool2d):
+
+    def forward(self, x):
+        if x.numel() == 0 and torch.__version__ <= '1.4':
+            out_shape = list(x.shape[:2])
+            for i, k, p, s, d in zip(x.shape[-2:], _pair(self.kernel_size),
+                                     _pair(self.padding), _pair(self.stride),
+                                     _pair(self.dilation)):
+                o = (i + 2 * p - (d * (k - 1) + 1)) / s + 1
+                o = math.ceil(o) if self.ceil_mode else math.floor(o)
+                out_shape.append(o)
+            empty = NewEmptyTensorOp.apply(x, out_shape)
+            return empty
+
+        return super().forward(x)
+
+
+class Linear(torch.nn.Linear):
+
+    def forward(self, x):
+        if x.numel() == 0:
+            out_shape = [x.shape[0], self.out_features]
+            empty = NewEmptyTensorOp.apply(x, out_shape)
+            if self.training:
+                # produce dummy gradient to avoid DDP warning.
+                dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
+                return empty + dummy
+            else:
+                return empty
+
+        return super().forward(x)