add cpp dcn extension

mmdet/ops/dcn/functions/deform_conv.py

+import torch
+from torch.autograd import Function
+from torch.nn.modules.utils import _pair
+
+from .. import deform_conv
+
+
+def deform_conv_function(input,
+                         offset,
+                         weight,
+                         stride=1,
+                         padding=0,
+                         dilation=1,
+                         deform_groups=1,
+                         im2col_step=64):
+    if input is not None and input.dim() != 4:
+        raise ValueError(
+            "Expected 4D tensor as input, got {}D tensor instead.".format(
+                input.dim()))
+
+    f = DeformConvFunction(
+        _pair(stride), _pair(padding), _pair(dilation), deform_groups,
+        im2col_step)
+    return f(input, offset, weight)
+
+
+class DeformConvFunction(Function):
+
+    def __init__(self,
+                 stride,
+                 padding,
+                 dilation,
+                 deformable_groups=1,
+                 im2col_step=64):
+        super(DeformConvFunction, self).__init__()
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.deformable_groups = deformable_groups
+        self.im2col_step = im2col_step
+
+    def forward(self, input, offset, weight):
+        self.save_for_backward(input, offset, weight)
+
+        output = input.new(*self._output_size(input, weight))
+
+        self.bufs_ = [input.new(), input.new()]  # columns, ones
+
+        if not input.is_cuda:
+            raise NotImplementedError
+        else:
+            if isinstance(input, torch.autograd.Variable):
+                if not isinstance(input.data, torch.cuda.FloatTensor):
+                    raise NotImplementedError
+            else:
+                if not isinstance(input, torch.cuda.FloatTensor):
+                    raise NotImplementedError
+
+            cur_im2col_step = min(self.im2col_step, input.shape[0])
+            assert (input.shape[0] %
+                    cur_im2col_step) == 0, 'im2col step must divide batchsize'
+            deform_conv.deform_conv_forward_cuda(
+                input, weight, offset, output, self.bufs_[0], self.bufs_[1],
+                weight.size(3), weight.size(2), self.stride[1], self.stride[0],
+                self.padding[1], self.padding[0], self.dilation[1],
+                self.dilation[0], self.deformable_groups, cur_im2col_step)
+        return output
+
+    def backward(self, grad_output):
+        input, offset, weight = self.saved_tensors
+
+        grad_input = grad_offset = grad_weight = None
+
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+        else:
+            if isinstance(grad_output, torch.autograd.Variable):
+                if not isinstance(grad_output.data, torch.cuda.FloatTensor):
+                    raise NotImplementedError
+            else:
+                if not isinstance(grad_output, torch.cuda.FloatTensor):
+                    raise NotImplementedError
+
+            cur_im2col_step = min(self.im2col_step, input.shape[0])
+            assert (input.shape[0] %
+                    cur_im2col_step) == 0, 'im2col step must divide batchsize'
+
+            if self.needs_input_grad[0] or self.needs_input_grad[1]:
+                grad_input = input.new(*input.size()).zero_()
+                grad_offset = offset.new(*offset.size()).zero_()
+                deform_conv.deform_conv_backward_input_cuda(
+                    input, offset, grad_output, grad_input,
+                    grad_offset, weight, self.bufs_[0], weight.size(3),
+                    weight.size(2), self.stride[1], self.stride[0],
+                    self.padding[1], self.padding[0], self.dilation[1],
+                    self.dilation[0], self.deformable_groups, cur_im2col_step)
+
+            if self.needs_input_grad[2]:
+                grad_weight = weight.new(*weight.size()).zero_()
+                deform_conv.deform_conv_backward_parameters_cuda(
+                    input, offset, grad_output,
+                    grad_weight, self.bufs_[0], self.bufs_[1], weight.size(3),
+                    weight.size(2), self.stride[1], self.stride[0],
+                    self.padding[1], self.padding[0], self.dilation[1],
+                    self.dilation[0], self.deformable_groups, 1,
+                    cur_im2col_step)
+
+        return grad_input, grad_offset, grad_weight
+
+    def _output_size(self, input, weight):
+        channels = weight.size(0)
+
+        output_size = (input.size(0), channels)
+        for d in range(input.dim() - 2):
+            in_size = input.size(d + 2)
+            pad = self.padding[d]
+            kernel = self.dilation[d] * (weight.size(d + 2) - 1) + 1
+            stride = self.stride[d]
+            output_size += ((in_size + (2 * pad) - kernel) // stride + 1, )
+        if not all(map(lambda s: s > 0, output_size)):
+            raise ValueError(
+                "convolution input is too small (output would be {})".format(
+                    'x'.join(map(str, output_size))))
+        return output_size

mmdet/ops/dcn/functions/modulated_dcn_func.py

+#!/usr/bin/env python
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import torch
+from torch.autograd import Function
+
+from .. import modulated_dcn as _backend
+
+
+class ModulatedDeformConvFunction(Function):
+
+    def __init__(self, stride, padding, dilation=1, deformable_groups=1):
+        super(ModulatedDeformConvFunction, self).__init__()
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.deformable_groups = deformable_groups
+
+    def forward(self, input, offset, mask, weight, bias):
+        if not input.is_cuda:
+            raise NotImplementedError
+        if weight.requires_grad or mask.requires_grad or offset.requires_grad \
+                or input.requires_grad:
+            self.save_for_backward(input, offset, mask, weight, bias)
+        output = input.new(*self._infer_shape(input, weight))
+        self._bufs = [input.new(), input.new()]
+        _backend.modulated_deform_conv_cuda_forward(
+            input, weight, bias, self._bufs[0], offset, mask, output,
+            self._bufs[1], weight.shape[2], weight.shape[3], self.stride,
+            self.stride, self.padding, self.padding, self.dilation,
+            self.dilation, self.deformable_groups)
+        return output
+
+    def backward(self, grad_output):
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+        input, offset, mask, weight, bias = self.saved_tensors
+        grad_input = input.new(*input.size()).zero_()
+        grad_offset = offset.new(*offset.size()).zero_()
+        grad_mask = mask.new(*mask.size()).zero_()
+        grad_weight = weight.new(*weight.size()).zero_()
+        grad_bias = bias.new(*bias.size()).zero_()
+        _backend.modulated_deform_conv_cuda_backward(
+            input, weight, bias, self._bufs[0], offset, mask, self._bufs[1],
+            grad_input, grad_weight, grad_bias, grad_offset, grad_mask,
+            grad_output, weight.shape[2], weight.shape[3], self.stride,
+            self.stride, self.padding, self.padding, self.dilation,
+            self.dilation, self.deformable_groups)
+
+        return grad_input, grad_offset, grad_mask, grad_weight, grad_bias
+
+    def _infer_shape(self, input, weight):
+        n = input.size(0)
+        channels_out = weight.size(0)
+        height, width = input.shape[2:4]
+        kernel_h, kernel_w = weight.shape[2:4]
+        height_out = (height + 2 * self.padding -
+                      (self.dilation * (kernel_h - 1) + 1)) // self.stride + 1
+        width_out = (width + 2 * self.padding -
+                     (self.dilation * (kernel_w - 1) + 1)) // self.stride + 1
+        return (n, channels_out, height_out, width_out)
+
+
+class DeformRoIPoolingFunction(Function):
+
+    def __init__(self,
+                 spatial_scale,
+                 pooled_size,
+                 output_dim,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0):
+        super(DeformRoIPoolingFunction, self).__init__()
+        self.spatial_scale = spatial_scale
+        self.pooled_size = pooled_size
+        self.output_dim = output_dim
+        self.no_trans = no_trans
+        self.group_size = group_size
+        self.part_size = pooled_size if part_size is None else part_size
+        self.sample_per_part = sample_per_part
+        self.trans_std = trans_std
+
+        assert self.trans_std >= 0.0 and self.trans_std <= 1.0
+
+    def forward(self, data, rois, offset):
+        if not data.is_cuda:
+            raise NotImplementedError
+
+        output = data.new(*self._infer_shape(data, rois))
+        output_count = data.new(*self._infer_shape(data, rois))
+        _backend.deform_psroi_pooling_cuda_forward(
+            data, rois, offset, output, output_count, self.no_trans,
+            self.spatial_scale, self.output_dim, self.group_size,
+            self.pooled_size, self.part_size, self.sample_per_part,
+            self.trans_std)
+
+        # if data.requires_grad or rois.requires_grad or offset.requires_grad:
+        #     self.save_for_backward(data, rois, offset, output_count)
+        self.data = data
+        self.rois = rois
+        self.offset = offset
+        self.output_count = output_count
+
+        return output
+
+    def backward(self, grad_output):
+        if not grad_output.is_cuda:
+            raise NotImplementedError
+
+        # data, rois, offset, output_count = self.saved_tensors
+        data = self.data
+        rois = self.rois
+        offset = self.offset
+        output_count = self.output_count
+        grad_input = data.new(*data.size()).zero_()
+        grad_offset = offset.new(*offset.size()).zero_()
+
+        _backend.deform_psroi_pooling_cuda_backward(
+            grad_output, data, rois, offset, output_count, grad_input,
+            grad_offset, self.no_trans, self.spatial_scale, self.output_dim,
+            self.group_size, self.pooled_size, self.part_size,
+            self.sample_per_part, self.trans_std)
+        return grad_input, torch.zeros(rois.shape).cuda(), grad_offset
+
+    def _infer_shape(self, data, rois):
+        # _, c, h, w = data.shape[:4]
+        # c = data.shape[1]
+        n = rois.shape[0]
+        return n, self.output_dim, self.pooled_size, self.pooled_size

mmdet/ops/dcn/modules/deform_conv.py

+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.module import Module
+from torch.nn.modules.utils import _pair
+
+from ..functions.deform_conv import deform_conv_function
+
+
+class DeformConv(Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 num_deformable_groups=1):
+        super(DeformConv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = _pair(stride)
+        self.padding = _pair(padding)
+        self.dilation = _pair(dilation)
+        self.num_deformable_groups = num_deformable_groups
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels, *self.kernel_size))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, offset):
+        return deform_conv_function(input, offset, self.weight, self.stride,
+                                    self.padding, self.dilation,
+                                    self.num_deformable_groups)

mmdet/ops/dcn/modules/modulated_dcn.py

+#!/usr/bin/env python
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+
+import torch
+from torch import nn
+from torch.nn.modules.utils import _pair
+
+from ..functions.modulated_dcn_func import DeformRoIPoolingFunction
+from ..functions.modulated_dcn_func import ModulatedDeformConvFunction
+
+
+class ModulatedDeformConv(nn.Module):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation=1,
+                 deformable_groups=1,
+                 no_bias=True):
+        super(ModulatedDeformConv, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = _pair(kernel_size)
+        self.stride = stride
+        self.padding = padding
+        self.dilation = dilation
+        self.deformable_groups = deformable_groups
+        self.no_bias = no_bias
+
+        self.weight = nn.Parameter(
+            torch.Tensor(out_channels, in_channels, *self.kernel_size))
+        self.bias = nn.Parameter(torch.zeros(out_channels))
+        self.reset_parameters()
+        if self.no_bias:
+            self.bias.requires_grad = False
+
+    def reset_parameters(self):
+        n = self.in_channels
+        for k in self.kernel_size:
+            n *= k
+        stdv = 1. / math.sqrt(n)
+        self.weight.data.uniform_(-stdv, stdv)
+        self.bias.data.zero_()
+
+    def forward(self, input, offset, mask):
+        func = ModulatedDeformConvFunction(
+            self.stride, self.padding, self.dilation, self.deformable_groups)
+        return func(input, offset, mask, self.weight, self.bias)
+
+
+class ModulatedDeformConvPack(ModulatedDeformConv):
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride,
+                 padding,
+                 dilation=1,
+                 deformable_groups=1,
+                 no_bias=False):
+        super(ModulatedDeformConvPack,
+              self).__init__(in_channels, out_channels, kernel_size, stride,
+                             padding, dilation, deformable_groups, no_bias)
+
+        self.conv_offset_mask = nn.Conv2d(
+            self.in_channels,
+            self.deformable_groups * 3 * self.kernel_size[0] *
+            self.kernel_size[1],
+            kernel_size=self.kernel_size,
+            stride=(self.stride, self.stride),
+            padding=(self.padding, self.padding),
+            bias=True)
+        self.init_offset()
+
+    def init_offset(self):
+        self.conv_offset_mask.weight.data.zero_()
+        self.conv_offset_mask.bias.data.zero_()
+
+    def forward(self, input):
+        out = self.conv_offset_mask(input)
+        o1, o2, mask = torch.chunk(out, 3, dim=1)
+        offset = torch.cat((o1, o2), dim=1)
+        mask = torch.sigmoid(mask)
+        func = ModulatedDeformConvFunction(
+            self.stride, self.padding, self.dilation, self.deformable_groups)
+        return func(input, offset, mask, self.weight, self.bias)
+
+
+class DeformRoIPooling(nn.Module):
+
+    def __init__(self,
+                 spatial_scale,
+                 pooled_size,
+                 output_dim,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0):
+        super(DeformRoIPooling, self).__init__()
+        self.spatial_scale = spatial_scale
+        self.pooled_size = pooled_size
+        self.out_size = pooled_size
+        self.output_dim = output_dim
+        self.no_trans = no_trans
+        self.group_size = group_size
+        self.part_size = pooled_size if part_size is None else part_size
+        self.sample_per_part = sample_per_part
+        self.trans_std = trans_std
+        self.func = DeformRoIPoolingFunction(
+            self.spatial_scale, self.pooled_size, self.output_dim,
+            self.no_trans, self.group_size, self.part_size,
+            self.sample_per_part, self.trans_std)
+
+    def forward(self, data, rois, offset):
+
+        if self.no_trans:
+            offset = data.new()
+        return self.func(data, rois, offset)
+
+
+class ModulatedDeformRoIPoolingPack(DeformRoIPooling):
+
+    def __init__(self,
+                 spatial_scale,
+                 pooled_size,
+                 output_dim,
+                 no_trans,
+                 group_size=1,
+                 part_size=None,
+                 sample_per_part=4,
+                 trans_std=.0,
+                 deform_fc_dim=1024):
+        super(ModulatedDeformRoIPoolingPack, self).__init__(
+            spatial_scale, pooled_size, output_dim, no_trans, group_size,
+            part_size, sample_per_part, trans_std)
+
+        self.deform_fc_dim = deform_fc_dim
+
+        if not no_trans:
+            self.func_offset = DeformRoIPoolingFunction(
+                self.spatial_scale, self.pooled_size, self.output_dim, True,
+                self.group_size, self.part_size, self.sample_per_part,
+                self.trans_std)
+            self.offset_fc = nn.Sequential(
+                nn.Linear(
+                    self.pooled_size * self.pooled_size * self.output_dim,
+                    self.deform_fc_dim), nn.ReLU(inplace=True),
+                nn.Linear(self.deform_fc_dim, self.deform_fc_dim),
+                nn.ReLU(inplace=True),
+                nn.Linear(self.deform_fc_dim,
+                          self.pooled_size * self.pooled_size * 2))
+            self.offset_fc[4].weight.data.zero_()
+            self.offset_fc[4].bias.data.zero_()
+            self.mask_fc = nn.Sequential(
+                nn.Linear(
+                    self.pooled_size * self.pooled_size * self.output_dim,
+                    self.deform_fc_dim), nn.ReLU(inplace=True),
+                nn.Linear(self.deform_fc_dim,
+                          self.pooled_size * self.pooled_size * 1),
+                nn.Sigmoid())
+            self.mask_fc[2].weight.data.zero_()
+            self.mask_fc[2].bias.data.zero_()
+
+    def forward(self, data, rois):
+        if self.no_trans:
+            offset = data.new()
+        else:
+            n = rois.shape[0]
+            offset = data.new()
+            x = self.func_offset(data, rois, offset)
+            offset = self.offset_fc(x.view(n, -1))
+            offset = offset.view(n, 2, self.pooled_size, self.pooled_size)
+            mask = self.mask_fc(x.view(n, -1))
+            mask = mask.view(n, 1, self.pooled_size, self.pooled_size)
+            feat = self.func(data, rois, offset) * mask
+            return feat
+        return self.func(data, rois, offset)

mmdet/ops/dcn/setup.py

+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='deform_conv',
+    ext_modules=[
+        CUDAExtension('deform_conv', [
+            'src/deform_conv_cuda.cpp',
+            'src/deform_conv_cuda_kernel.cu',
+        ]),
+    ],
+    cmdclass={'build_ext': BuildExtension})

mmdet/ops/dcn/setup_modulated.py

+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name='modulated_deform_conv',
+    ext_modules=[
+        CUDAExtension('modulated_dcn', [
+            'src/modulated_dcn_cuda.cpp',
+            'src/modulated_deform_im2col_cuda.cu',
+            'src/deform_psroi_pooling_cuda.cu'
+        ]),
+    ],
+    cmdclass={'build_ext': BuildExtension})

mmdet/ops/dcn/src/deform_conv_cuda.cpp

+#include <torch/torch.h>
+
+#include <cmath>
+#include <vector>
+
+void deformable_im2col(const at::Tensor data_im,
+                       const at::Tensor data_offset, const int channels,
+                       const int height, const int width, const int ksize_h,
+                       const int ksize_w, const int pad_h, const int pad_w,
+                       const int stride_h, const int stride_w,
+                       const int dilation_h, const int dilation_w,
+                       const int parallel_imgs,
+                       const int deformable_group, at::Tensor data_col);
+
+void deformable_col2im(const at::Tensor data_col,
+                       const at::Tensor data_offset, const int channels,
+                       const int height, const int width, const int ksize_h,
+                       const int ksize_w, const int pad_h, const int pad_w,
+                       const int stride_h, const int stride_w,
+                       const int dilation_h, const int dilation_w,
+                       const int parallel_imgs,
+                       const int deformable_group, at::Tensor grad_im);
+
+void deformable_col2im_coord(const at::Tensor data_col,
+                             const at::Tensor data_im, const at::Tensor data_offset,
+                             const int channels, const int height,
+                             const int width, const int ksize_h,
+                             const int ksize_w, const int pad_h,
+                             const int pad_w, const int stride_h,
+                             const int stride_w, const int dilation_h,
+                             const int dilation_w, const int parallel_imgs,
+                             const int deformable_group, at::Tensor grad_offset);
+
+void shape_check(at::Tensor input, at::Tensor offset,
+                 at::Tensor *gradOutput, at::Tensor weight, int kH, int kW,
+                 int dH, int dW, int padH, int padW, int dilationH,
+                 int dilationW, int deformable_group)
+{
+
+    //  AT_CHECK(weight->nDimension == 4, 5,
+    //             "4D weight tensor (nOutputPlane,nInputPlane,kH,kW) expected, "
+    //             "but got: %s",
+    //             weight->nDimension);
+    AT_CHECK(weight.ndimension() == 4,
+             "4D weight tensor (nOutputPlane,nInputPlane,kH,kW) expected, "
+             "but got: %s",
+             weight.ndimension());
+
+    AT_CHECK(weight.is_contiguous(),
+             "weight tensor has to be contiguous");
+
+    AT_CHECK(kW > 0 && kH > 0,
+             "kernel size should be greater than zero, but got kH: %d kW: %d",
+             kH, kW);
+
+    //  AT_CHECK((weight->size[2] == kH && weight->size[3] == kW), 9,
+    //             "kernel size should be consistent with weight, ",
+    //             "but got kH: %d kW: %d weight.size(2): %d, weight.size(3): %d", kH,
+    //             kW, weight->size[2], weight->size[3]);
+    AT_CHECK((weight.size(2) == kH &&
+              weight.size(3) == kW),
+             "kernel size should be consistent with weight, ",
+             "but got kH: %d kW: %d weight.size(2): %d, weight.size(3): %d", kH,
+             kW, weight.size(2), weight.size(3));
+
+    AT_CHECK(dW > 0 && dH > 0,
+             "stride should be greater than zero, but got dH: %d dW: %d", dH, dW);
+
+    AT_CHECK(dilationW > 0 && dilationH > 0,
+             "dilation should be greater than 0, but got dilationH: %d dilationW: %d",
+             dilationH, dilationW);
+
+    //  int ndim = input->nDimension;
+    int ndim = input.ndimension();
+    int dimf = 0;
+    int dimh = 1;
+    int dimw = 2;
+
+    if (ndim == 4)
+    {
+        dimf++;
+        dimh++;
+        dimw++;
+    }
+
+    AT_CHECK(ndim == 3 || ndim == 4,
+             "3D or 4D input tensor expected but got: %s", ndim);
+
+    //  long nInputPlane = weight->size[1];
+    //  long inputHeight = input->size[dimh];
+    //  long inputWidth = input->size[dimw];
+    //  long nOutputPlane = weight->size[0];
+    long nInputPlane = weight.size(1);
+    long inputHeight = input.size(dimh);
+    long inputWidth = input.size(dimw);
+    long nOutputPlane = weight.size(0);
+    long outputHeight = (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+    long outputWidth = (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+
+    AT_CHECK(nInputPlane % deformable_group == 0,
+             "input channels must divide deformable group size");
+
+    if (outputWidth < 1 || outputHeight < 1)
+        AT_ERROR(
+            "Given input size: (%ld x %ld x %ld). "
+            "Calculated output size: (%ld x %ld x %ld). Output size is too small",
+            nInputPlane, inputHeight, inputWidth, nOutputPlane, outputHeight,
+            outputWidth);
+
+    AT_CHECK(input.size(1) == nInputPlane,
+             "invalid number of input planes, expected: %d, but got: %d",
+             nInputPlane, input.size(1));
+
+    AT_CHECK((inputHeight >= kH && inputWidth >= kW),
+             "input image is smaller than kernel");
+
+    //  AT_CHECK(
+    //      (offset->size[2] == outputHeight && offset->size[3] == outputWidth), 3,
+    //      "invalid spatial size of offset, expected height: %d width: %d, but got height: %d width: %d", outputHeight, outputWidth,
+    //      offset->size[2], offset->size[3]);
+    AT_CHECK(
+        (offset.size(2) == outputHeight && offset.size(3) == outputWidth),
+        "invalid spatial size of offset, expected height: %d width: %d, but got height: %d width: %d",
+        outputHeight, outputWidth, offset.size(2), offset.size(3));
+
+    AT_CHECK((offset.size(1) == deformable_group * 2 * kH * kW),
+             "invalid number of channels of offset");
+
+    if (gradOutput != NULL)
+    {
+        AT_CHECK(gradOutput->size(dimf) == nOutputPlane,
+                 "invalid number of gradOutput planes, expected: %d, but got: %d",
+                 nOutputPlane, gradOutput->size(dimf));
+
+        AT_CHECK((gradOutput->size(dimh) == outputHeight &&
+                  gradOutput->size(dimw) == outputWidth),
+                 "invalid size of gradOutput, expected height: %d width: %d , but got height: %d width: %d",
+                 outputHeight, outputWidth, gradOutput->size(dimh), gradOutput->size(dimw));
+    }
+}
+
+int deform_conv_forward_cuda(at::Tensor input, at::Tensor weight,
+                             at::Tensor offset, at::Tensor output,
+                             at::Tensor columns, at::Tensor ones, int kW,
+                             int kH, int dW, int dH, int padW, int padH,
+                             int dilationW, int dilationH,
+                             int deformable_group, int im2col_step)
+{
+
+    // todo: resize columns to include im2col: done
+    // todo: add im2col_step as input
+    // todo: add new output buffer and transpose it to output (or directly transpose output)
+    // todo: possibly change data indexing because of parallel_imgs
+
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 6, input, weight, offset,
+    //                                        output, columns, ones));
+
+    shape_check(input, offset, NULL, weight, kH, kW, dH, dW, padH, padW,
+                dilationH, dilationW, deformable_group);
+
+    input = input.contiguous();
+    offset = offset.contiguous();
+    weight = weight.contiguous();
+
+    int batch = 1;
+    if (input.ndimension() == 3)
+    {
+        // Force batch
+        batch = 0;
+        input.unsqueeze_(0);
+        offset.unsqueeze_(0);
+    }
+
+    // todo: assert batchsize dividable by im2col_step
+
+    long batchSize = input.size(0);
+    long nInputPlane = input.size(1);
+    long inputHeight = input.size(2);
+    long inputWidth = input.size(3);
+
+    long nOutputPlane = weight.size(0);
+
+    long outputWidth = (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+    long outputHeight = (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+    AT_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
+
+    // bias = bias ? THCudaTensor_newContiguous(state, bias) : bias;
+
+    output = output.view({batchSize / im2col_step, im2col_step, nOutputPlane, outputHeight, outputWidth});
+    columns = at::zeros({nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth}, input.type());
+
+    if (ones.ndimension() != 2 || ones.size(0) * ones.size(1) < outputHeight * outputWidth)
+    {
+        ones = at::ones({outputHeight, outputWidth}, input.type());
+    }
+
+    input = input.view({batchSize / im2col_step, im2col_step, nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({batchSize / im2col_step, im2col_step,
+                          deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+    at::Tensor output_buffer = at::zeros({batchSize / im2col_step, nOutputPlane, im2col_step * outputHeight, outputWidth}, output.type());
+
+    for (int elt = 0; elt < batchSize / im2col_step; elt++)
+    {
+        deformable_im2col(
+            input[elt], offset[elt], nInputPlane, inputHeight,
+            inputWidth, kH, kW, padH, padW, dH, dW, dilationH, dilationW,
+            im2col_step, deformable_group, columns);
+
+        output_buffer[elt] =
+            output_buffer[elt].flatten(1).addmm_(weight.flatten(1), columns).view_as(output_buffer[elt]);
+    }
+
+    // the reason I use seemingly redundant output_buffer is that THCudaTensor API handles successive transpose and resize poorly
+    output_buffer = output_buffer.view(
+        {batchSize / im2col_step, nOutputPlane, im2col_step, outputHeight, outputWidth});
+    output_buffer.transpose_(1, 2);
+    output.copy_(output_buffer);
+    output = output.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+    input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+    if (batch == 0)
+    {
+        output = output.view({nOutputPlane, outputHeight, outputWidth});
+        input = input.view({nInputPlane, inputHeight, inputWidth});
+        offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
+    }
+
+    return 1;
+}
+
+int deform_conv_backward_input_cuda(
+    at::Tensor input, at::Tensor offset, at::Tensor gradOutput,
+    at::Tensor gradInput, at::Tensor gradOffset, at::Tensor weight,
+    at::Tensor columns, int kW, int kH, int dW, int dH, int padW, int padH,
+    int dilationW, int dilationH, int deformable_group, int im2col_step)
+{
+
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 6, input, gradOutput, weight,
+    //                                        offset, columns, gradInput));
+
+    shape_check(input, offset, &gradOutput, weight, kH, kW, dH, dW, padH,
+                padW, dilationH, dilationW, deformable_group);
+
+    input = input.contiguous();
+    offset = offset.contiguous();
+    gradOutput = gradOutput.contiguous();
+    weight = weight.contiguous();
+
+    int batch = 1;
+
+    if (input.ndimension() == 3)
+    {
+        // Force batch
+        batch = 0;
+        input = input.view({1, input.size(0), input.size(1), input.size(2)});
+        offset = offset.view({1, offset.size(0), offset.size(1), offset.size(2)});
+        gradOutput = gradOutput.view({1, gradOutput.size(0), gradOutput.size(1), gradOutput.size(2)});
+    }
+
+    long batchSize = input.size(0);
+    long nInputPlane = input.size(1);
+    long inputHeight = input.size(2);
+    long inputWidth = input.size(3);
+
+    long nOutputPlane = weight.size(0);
+
+    long outputWidth = (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+    long outputHeight = (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+    AT_CHECK((offset.size(0) == batchSize), 3, "invalid batch size of offset");
+    gradInput = gradInput.view({batchSize, nInputPlane, inputHeight, inputWidth});
+    columns = at::zeros({nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth}, input.type());
+
+    // change order of grad output
+    gradOutput = gradOutput.view(
+        {batchSize / im2col_step, im2col_step, nOutputPlane, outputHeight, outputWidth});
+    gradOutput.transpose_(1, 2);
+
+    at::Tensor gradOutputBuffer = at::zeros_like(gradOutput);
+    gradOutputBuffer = gradOutputBuffer.view(
+        {batchSize / im2col_step, nOutputPlane, im2col_step, outputHeight, outputWidth});
+    gradOutputBuffer.copy_(gradOutput);
+    gradOutputBuffer = gradOutputBuffer.view(
+        {batchSize / im2col_step, nOutputPlane, im2col_step * outputHeight, outputWidth});
+
+    gradOutput.transpose_(1, 2);
+    gradOutput = gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+    gradInput = gradInput.view(
+        {batchSize / im2col_step, im2col_step, nInputPlane, inputHeight, inputWidth});
+    input = input.view({batchSize / im2col_step, im2col_step, nInputPlane, inputHeight, inputWidth});
+    gradOffset = gradOffset.view({batchSize / im2col_step, im2col_step,
+                                  deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+    offset = offset.view({batchSize / im2col_step, im2col_step,
+                          deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+    for (int elt = 0; elt < batchSize / im2col_step; elt++)
+    {
+        columns = columns.addmm_(weight.flatten(1).transpose(0, 1), gradOutputBuffer[elt].flatten(1));
+
+        deformable_col2im_coord(
+            columns, input[elt], offset[elt],
+            nInputPlane, inputHeight, inputWidth, kH, kW, padH, padW, dH, dW,
+            dilationH, dilationW, im2col_step, deformable_group, gradOffset[elt]);
+
+        deformable_col2im(
+            columns, offset[elt], nInputPlane, inputHeight,
+            inputWidth, kH, kW, padH, padW, dH, dW, dilationH, dilationW, im2col_step,
+            deformable_group, gradInput[elt]);
+    }
+
+    gradInput = gradInput.view({batchSize, nInputPlane, inputHeight, inputWidth});
+    input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+    gradOffset = gradOffset.view({batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+    offset = offset.view({batchSize, deformable_group * 2 * kH * kW, outputHeight, outputWidth});
+
+    if (batch == 0)
+    {
+        gradOutput = gradOutput.view({nOutputPlane, outputHeight, outputWidth});
+        input = input.view({nInputPlane, inputHeight, inputWidth});
+        gradInput = gradInput.view({nInputPlane, inputHeight, inputWidth});
+        offset = offset.view({offset.size(1), offset.size(2), offset.size(3)});
+        gradOffset = gradOffset.view({offset.size(1), offset.size(2), offset.size(3)});
+    }
+
+    return 1;
+}
+
+int deform_conv_backward_parameters_cuda(
+    at::Tensor input, at::Tensor offset, at::Tensor gradOutput,
+    at::Tensor gradWeight, // at::Tensor gradBias,
+    at::Tensor columns, at::Tensor ones, int kW, int kH, int dW, int dH,
+    int padW, int padH, int dilationW, int dilationH, int deformable_group,
+    float scale, int im2col_step)
+{
+
+    // todo: transpose and reshape outGrad
+    // todo: reshape columns
+    // todo: add im2col_step as input
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 5, input, offset, gradOutput,
+    //                                        gradWeight, columns));
+
+    shape_check(input, offset, &gradOutput, gradWeight, kH, kW, dH, dW,
+                padH, padW, dilationH, dilationW, deformable_group);
+
+    input = input.contiguous();
+    offset = offset.contiguous();
+    gradOutput = gradOutput.contiguous();
+
+    int batch = 1;
+
+    if (input.ndimension() == 3)
+    {
+        // Force batch
+        batch = 0;
+        input = input.view(at::IntList({1, input.size(0), input.size(1), input.size(2)}));
+        gradOutput = gradOutput.view({1, gradOutput.size(0),
+                                      gradOutput.size(1), gradOutput.size(2)});
+    }
+
+    long batchSize = input.size(0);
+    long nInputPlane = input.size(1);
+    long inputHeight = input.size(2);
+    long inputWidth = input.size(3);
+
+    long nOutputPlane = gradWeight.size(0);
+
+    long outputWidth = (inputWidth + 2 * padW - (dilationW * (kW - 1) + 1)) / dW + 1;
+    long outputHeight = (inputHeight + 2 * padH - (dilationH * (kH - 1) + 1)) / dH + 1;
+
+    AT_CHECK((offset.size(0) == batchSize), "invalid batch size of offset");
+
+    columns = at::zeros({nInputPlane * kW * kH, im2col_step * outputHeight * outputWidth}, input.type());
+
+    gradOutput = gradOutput.view(
+        {batchSize / im2col_step, im2col_step, nOutputPlane, outputHeight, outputWidth});
+    gradOutput.transpose_(1, 2);
+
+    at::Tensor gradOutputBuffer = at::zeros_like(gradOutput);
+    gradOutputBuffer = gradOutputBuffer.view(
+        {batchSize / im2col_step, nOutputPlane, im2col_step, outputHeight, outputWidth});
+    gradOutputBuffer.copy_(gradOutput);
+    gradOutputBuffer = gradOutputBuffer.view(
+        {batchSize / im2col_step, nOutputPlane, im2col_step * outputHeight, outputWidth});
+
+    gradOutput.transpose_(1, 2);
+    gradOutput = gradOutput.view({batchSize, nOutputPlane, outputHeight, outputWidth});
+
+    input = input.view({batchSize / im2col_step, im2col_step, nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({batchSize / im2col_step, im2col_step,
+                          deformable_group * 2 * kH * kW,
+                          outputHeight, outputWidth});
+
+    for (int elt = 0; elt < batchSize / im2col_step; elt++)
+    {
+        deformable_im2col(
+            input[elt], offset[elt], nInputPlane, inputHeight,
+            inputWidth, kH, kW, padH, padW, dH, dW, dilationH, dilationW,
+            im2col_step, deformable_group, columns);
+
+        gradWeight.copy_(gradWeight.flatten(1).addmm_(
+            gradOutputBuffer[elt].flatten(1), columns.transpose(1, 0), 1.0, scale).view_as(gradWeight));
+    }
+
+    input = input.view({batchSize, nInputPlane, inputHeight, inputWidth});
+    offset = offset.view({batchSize, deformable_group * 2 * kH * kW,
+                          outputHeight, outputWidth});
+
+    if (batch == 0)
+    {
+        gradOutput = gradOutput.view({nOutputPlane, outputHeight, outputWidth});
+        input = input.view({nInputPlane, inputHeight, inputWidth});
+    }
+
+    return 1;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("deform_conv_forward_cuda", &deform_conv_forward_cuda, "deform forward (CUDA)");
+    m.def("deform_conv_backward_input_cuda", &deform_conv_backward_input_cuda,
+          "deform_conv_backward_input (CUDA)");
+    m.def("deform_conv_backward_parameters_cuda", &deform_conv_backward_parameters_cuda,
+          "deform_conv_backward_parameters (CUDA)");
+}

mmdet/ops/dcn/src/deform_psroi_pooling_cuda.cu

+/*!
+ * Copyright (c) 2017 Microsoft
+ * Licensed under The MIT License [see LICENSE for details]
+ * \file deformable_psroi_pooling.cu
+ * \brief
+ * \author Yi Li, Guodong Zhang, Jifeng Dai
+*/
+/***************** Adapted by Charles Shang *********************/
+#include <ATen/ATen.h>
+#include <THC/THCAtomics.cuh>
+#include <stdio.h>
+#include <math.h>
+#include <algorithm>
+
+using namespace at;
+
+#define CUDA_KERNEL_LOOP(i, n)                        \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; \
+       i < (n);                                       \
+       i += blockDim.x * gridDim.x)
+
+const int CUDA_NUM_THREADS = 1024;
+inline int GET_BLOCKS(const int N)
+{
+  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
+}
+
+template <typename scalar_t>
+__device__ scalar_t bilinear_interp(
+    const scalar_t *data,
+    const scalar_t x,
+    const scalar_t y,
+    const int width,
+    const int height)
+{
+  int x1 = floor(x);
+  int x2 = ceil(x);
+  int y1 = floor(y);
+  int y2 = ceil(y);
+  scalar_t dist_x = (scalar_t)(x - x1);
+  scalar_t dist_y = (scalar_t)(y - y1);
+  scalar_t value11 = data[y1 * width + x1];
+  scalar_t value12 = data[y2 * width + x1];
+  scalar_t value21 = data[y1 * width + x2];
+  scalar_t value22 = data[y2 * width + x2];
+  scalar_t value = (1 - dist_x) * (1 - dist_y) * value11 + (1 - dist_x) * dist_y * value12 + dist_x * (1 - dist_y) * value21 + dist_x * dist_y * value22;
+  return value;
+}
+
+template <typename scalar_t>
+__global__ void DeformablePSROIPoolForwardKernel(
+    const int count,
+    const scalar_t *bottom_data,
+    const scalar_t spatial_scale,
+    const int channels,
+    const int height, const int width,
+    const int pooled_height, const int pooled_width,
+    const scalar_t *bottom_rois, const scalar_t *bottom_trans,
+    const int no_trans,
+    const scalar_t trans_std,
+    const int sample_per_part,
+    const int output_dim,
+    const int group_size,
+    const int part_size,
+    const int num_classes,
+    const int channels_each_class,
+    scalar_t *top_data,
+    scalar_t *top_count)
+{
+  CUDA_KERNEL_LOOP(index, count)
+  {
+    // The output is in order (n, ctop, ph, pw)
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int ctop = (index / pooled_width / pooled_height) % output_dim;
+    int n = index / pooled_width / pooled_height / output_dim;
+
+    // [start, end) interval for spatial sampling
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 5;
+    int roi_batch_ind = offset_bottom_rois[0];
+    scalar_t roi_start_w = (scalar_t)(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
+    scalar_t roi_start_h = (scalar_t)(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
+    scalar_t roi_end_w = (scalar_t)(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
+    scalar_t roi_end_h = (scalar_t)(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
+
+    // Force too small ROIs to be 1x1
+    scalar_t roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
+    scalar_t roi_height = max(roi_end_h - roi_start_h, 0.1);
+
+    // Compute w and h at bottom
+    scalar_t bin_size_h = roi_height / (scalar_t)(pooled_height);
+    scalar_t bin_size_w = roi_width / (scalar_t)(pooled_width);
+
+    scalar_t sub_bin_size_h = bin_size_h / (scalar_t)(sample_per_part);
+    scalar_t sub_bin_size_w = bin_size_w / (scalar_t)(sample_per_part);
+
+    int part_h = floor((scalar_t)(ph) / pooled_height * part_size);
+    int part_w = floor((scalar_t)(pw) / pooled_width * part_size);
+    int class_id = ctop / channels_each_class;
+    scalar_t trans_x = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+    scalar_t trans_y = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+
+    scalar_t wstart = (scalar_t)(pw)*bin_size_w + roi_start_w;
+    wstart += trans_x * roi_width;
+    scalar_t hstart = (scalar_t)(ph)*bin_size_h + roi_start_h;
+    hstart += trans_y * roi_height;
+
+    scalar_t sum = 0;
+    int count = 0;
+    int gw = floor((scalar_t)(pw)*group_size / pooled_width);
+    int gh = floor((scalar_t)(ph)*group_size / pooled_height);
+    gw = min(max(gw, 0), group_size - 1);
+    gh = min(max(gh, 0), group_size - 1);
+
+    const scalar_t *offset_bottom_data = bottom_data + (roi_batch_ind * channels) * height * width;
+    for (int ih = 0; ih < sample_per_part; ih++)
+    {
+      for (int iw = 0; iw < sample_per_part; iw++)
+      {
+        scalar_t w = wstart + iw * sub_bin_size_w;
+        scalar_t h = hstart + ih * sub_bin_size_h;
+        // bilinear interpolation
+        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
+        {
+          continue;
+        }
+        w = min(max(w, 0.), width - 1.);
+        h = min(max(h, 0.), height - 1.);
+        int c = (ctop * group_size + gh) * group_size + gw;
+        scalar_t val = bilinear_interp(offset_bottom_data + c * height * width, w, h, width, height);
+        sum += val;
+        count++;
+      }
+    }
+    top_data[index] = count == 0 ? (scalar_t)(0) : sum / count;
+    top_count[index] = count;
+  }
+}
+
+template <typename scalar_t>
+__global__ void DeformablePSROIPoolBackwardAccKernel(
+    const int count,
+    const scalar_t *top_diff,
+    const scalar_t *top_count,
+    const int num_rois,
+    const scalar_t spatial_scale,
+    const int channels,
+    const int height, const int width,
+    const int pooled_height, const int pooled_width,
+    const int output_dim,
+    scalar_t *bottom_data_diff, scalar_t *bottom_trans_diff,
+    const scalar_t *bottom_data,
+    const scalar_t *bottom_rois,
+    const scalar_t *bottom_trans,
+    const int no_trans,
+    const scalar_t trans_std,
+    const int sample_per_part,
+    const int group_size,
+    const int part_size,
+    const int num_classes,
+    const int channels_each_class)
+{
+  CUDA_KERNEL_LOOP(index, count)
+  {
+    // The output is in order (n, ctop, ph, pw)
+    int pw = index % pooled_width;
+    int ph = (index / pooled_width) % pooled_height;
+    int ctop = (index / pooled_width / pooled_height) % output_dim;
+    int n = index / pooled_width / pooled_height / output_dim;
+
+    // [start, end) interval for spatial sampling
+    const scalar_t *offset_bottom_rois = bottom_rois + n * 5;
+    int roi_batch_ind = offset_bottom_rois[0];
+    scalar_t roi_start_w = (scalar_t)(round(offset_bottom_rois[1])) * spatial_scale - 0.5;
+    scalar_t roi_start_h = (scalar_t)(round(offset_bottom_rois[2])) * spatial_scale - 0.5;
+    scalar_t roi_end_w = (scalar_t)(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;
+    scalar_t roi_end_h = (scalar_t)(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;
+
+    // Force too small ROIs to be 1x1
+    scalar_t roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0
+    scalar_t roi_height = max(roi_end_h - roi_start_h, 0.1);
+
+    // Compute w and h at bottom
+    scalar_t bin_size_h = roi_height / (scalar_t)(pooled_height);
+    scalar_t bin_size_w = roi_width / (scalar_t)(pooled_width);
+
+    scalar_t sub_bin_size_h = bin_size_h / (scalar_t)(sample_per_part);
+    scalar_t sub_bin_size_w = bin_size_w / (scalar_t)(sample_per_part);
+
+    int part_h = floor((scalar_t)(ph) / pooled_height * part_size);
+    int part_w = floor((scalar_t)(pw) / pooled_width * part_size);
+    int class_id = ctop / channels_each_class;
+    scalar_t trans_x = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+    scalar_t trans_y = no_trans ? (scalar_t)(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * (scalar_t)trans_std;
+
+    scalar_t wstart = (scalar_t)(pw)*bin_size_w + roi_start_w;
+    wstart += trans_x * roi_width;
+    scalar_t hstart = (scalar_t)(ph)*bin_size_h + roi_start_h;
+    hstart += trans_y * roi_height;
+
+    if (top_count[index] <= 0)
+    {
+      continue;
+    }
+    scalar_t diff_val = top_diff[index] / top_count[index];
+    const scalar_t *offset_bottom_data = bottom_data + roi_batch_ind * channels * height * width;
+    scalar_t *offset_bottom_data_diff = bottom_data_diff + roi_batch_ind * channels * height * width;
+    int gw = floor((scalar_t)(pw)*group_size / pooled_width);
+    int gh = floor((scalar_t)(ph)*group_size / pooled_height);
+    gw = min(max(gw, 0), group_size - 1);
+    gh = min(max(gh, 0), group_size - 1);
+
+    for (int ih = 0; ih < sample_per_part; ih++)
+    {
+      for (int iw = 0; iw < sample_per_part; iw++)
+      {
+        scalar_t w = wstart + iw * sub_bin_size_w;
+        scalar_t h = hstart + ih * sub_bin_size_h;
+        // bilinear interpolation
+        if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)
+        {
+          continue;
+        }
+        w = min(max(w, 0.), width - 1.);
+        h = min(max(h, 0.), height - 1.);
+        int c = (ctop * group_size + gh) * group_size + gw;
+        // backward on feature
+        int x0 = floor(w);
+        int x1 = ceil(w);
+        int y0 = floor(h);
+        int y1 = ceil(h);
+        scalar_t dist_x = w - x0, dist_y = h - y0;
+        scalar_t q00 = (1 - dist_x) * (1 - dist_y);
+        scalar_t q01 = (1 - dist_x) * dist_y;
+        scalar_t q10 = dist_x * (1 - dist_y);
+        scalar_t q11 = dist_x * dist_y;
+        int bottom_index_base = c * height * width;
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x0, q00 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x0, q01 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x1, q10 * diff_val);
+        atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x1, q11 * diff_val);
+
+        if (no_trans)
+        {
+          continue;
+        }
+        scalar_t U00 = offset_bottom_data[bottom_index_base + y0 * width + x0];
+        scalar_t U01 = offset_bottom_data[bottom_index_base + y1 * width + x0];
+        scalar_t U10 = offset_bottom_data[bottom_index_base + y0 * width + x1];
+        scalar_t U11 = offset_bottom_data[bottom_index_base + y1 * width + x1];
+        scalar_t diff_x = (U11 * dist_y + U10 * (1 - dist_y) - U01 * dist_y - U00 * (1 - dist_y)) * trans_std * diff_val;
+        diff_x *= roi_width;
+        scalar_t diff_y = (U11 * dist_x + U01 * (1 - dist_x) - U10 * dist_x - U00 * (1 - dist_x)) * trans_std * diff_val;
+        diff_y *= roi_height;
+
+        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w, diff_x);
+        atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w, diff_y);
+      }
+    }
+  }
+}
+
+void DeformablePSROIPoolForward(const at::Tensor data,
+                                const at::Tensor bbox,
+                                const at::Tensor trans,
+                                at::Tensor out,
+                                at::Tensor top_count,
+                                const int batch,
+                                const int channels,
+                                const int height,
+                                const int width,
+                                const int num_bbox,
+                                const int channels_trans,
+                                const int no_trans,
+                                const float spatial_scale,
+                                const int output_dim,
+                                const int group_size,
+                                const int pooled_size,
+                                const int part_size,
+                                const int sample_per_part,
+                                const float trans_std)
+{
+  const int pooled_height = pooled_size;
+  const int pooled_width = pooled_size;
+  const int count = num_bbox * output_dim * pooled_height * pooled_width;
+  const int num_classes = no_trans ? 1 : channels_trans / 2;
+  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+    data.type(), "deformable_psroi_pool_forward", ([&] {
+    const scalar_t *bottom_data = data.data<scalar_t>();
+    const scalar_t *bottom_rois = bbox.data<scalar_t>();
+    const scalar_t *bottom_trans = no_trans ? NULL : trans.data<scalar_t>();
+    scalar_t *top_data = out.data<scalar_t>();
+    scalar_t *top_count_data = top_count.data<scalar_t>();
+
+    DeformablePSROIPoolForwardKernel<<<GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
+      count, bottom_data, (scalar_t)spatial_scale, channels, height, width, pooled_height, pooled_width,
+      bottom_rois, bottom_trans, no_trans, (scalar_t)trans_std, sample_per_part, output_dim,
+      group_size, part_size, num_classes, channels_each_class, top_data, top_count_data);
+  }));
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in DeformablePSROIPoolForward: %s\n", cudaGetErrorString(err));
+  }
+}
+
+void DeformablePSROIPoolBackwardAcc(const at::Tensor out_grad,
+                                    const at::Tensor data,
+                                    const at::Tensor bbox,
+                                    const at::Tensor trans,
+                                    const at::Tensor top_count,
+                                    at::Tensor in_grad,
+                                    at::Tensor trans_grad,
+                                    const int batch,
+                                    const int channels,
+                                    const int height,
+                                    const int width,
+                                    const int num_bbox,
+                                    const int channels_trans,
+                                    const int no_trans,
+                                    const float spatial_scale,
+                                    const int output_dim,
+                                    const int group_size,
+                                    const int pooled_size,
+                                    const int part_size,
+                                    const int sample_per_part,
+                                    const float trans_std)
+{
+  // LOG(INFO) << "DeformablePSROIPoolBackward";
+  const int num_rois = num_bbox;
+  const int pooled_height = pooled_size;
+  const int pooled_width = pooled_size;
+  const int count = num_bbox * output_dim * pooled_height * pooled_width;
+  const int num_classes = no_trans ? 1 : channels_trans / 2;
+  const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      out_grad.type(), "deformable_psroi_pool_backward_acc", ([&] {
+        const scalar_t *top_diff = out_grad.data<scalar_t>();
+        const scalar_t *bottom_data = data.data<scalar_t>();
+        const scalar_t *bottom_rois = bbox.data<scalar_t>();
+        const scalar_t *bottom_trans = no_trans ? NULL : trans.data<scalar_t>();
+        scalar_t *bottom_data_diff = in_grad.data<scalar_t>();
+        scalar_t *bottom_trans_diff = no_trans ? NULL : trans_grad.data<scalar_t>();
+        const scalar_t *top_count_data = top_count.data<scalar_t>();
+
+        DeformablePSROIPoolBackwardAccKernel<<<GET_BLOCKS(count), CUDA_NUM_THREADS>>>(
+            count, top_diff, top_count_data, num_rois, (scalar_t)spatial_scale, channels, height, width,
+            pooled_height, pooled_width, output_dim, bottom_data_diff, bottom_trans_diff,
+            bottom_data, bottom_rois, bottom_trans, no_trans, (scalar_t)trans_std, sample_per_part,
+            group_size, part_size, num_classes, channels_each_class);
+      }));
+
+
+  cudaError_t err = cudaGetLastError();
+  if (err != cudaSuccess)
+  {
+    printf("error in DeformablePSROIPoolForward: %s\n", cudaGetErrorString(err));
+  }
+}
\ No newline at end of file

mmdet/ops/dcn/src/modulated_dcn_cuda.cpp

+#include <torch/torch.h>
+
+#include <cmath>
+#include <vector>
+
+// author: Charles Shang
+// https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu
+
+
+void modulated_deformable_im2col_cuda(const at::Tensor data_im, const at::Tensor data_offset,
+                                      const at::Tensor data_mask, const int batch_size, const int channels,
+                                      const int height_im, const int width_im, const int height_col,
+                                      const int width_col, const int kernel_h, const int kenerl_w,
+                                      const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+                                      const int dilation_h, const int dilation_w,
+                                      const int deformable_group, at::Tensor data_col);
+
+void modulated_deformable_col2im_cuda(const at::Tensor data_col, const at::Tensor data_offset,
+                                      const at::Tensor data_mask, const int batch_size, const int channels,
+                                      const int height_im, const int width_im, const int height_col,
+                                      const int width_col, const int kernel_h, const int kenerl_w,
+                                      const int pad_h, const int pad_w, const int stride_h, const int stride_w,
+                                      const int dilation_h, const int dilation_w,
+                                      const int deformable_group, at::Tensor grad_im);
+
+void modulated_deformable_col2im_coord_cuda(const at::Tensor data_col, const at::Tensor data_im,
+                                            const at::Tensor data_offset, const at::Tensor data_mask,
+                                            const int batch_size, const int channels, const int height_im,
+                                            const int width_im, const int height_col, const int width_col,
+                                            const int kernel_h, const int kenerl_w, const int pad_h,
+                                            const int pad_w, const int stride_h, const int stride_w,
+                                            const int dilation_h, const int dilation_w,
+                                            const int deformable_group, at::Tensor grad_offset,
+                                            at::Tensor grad_mask);
+
+
+void DeformablePSROIPoolForward(const at::Tensor data,
+                                const at::Tensor bbox,
+                                const at::Tensor trans,
+                                at::Tensor out,
+                                at::Tensor top_count,
+                                const int batch,
+                                const int channels,
+                                const int height,
+                                const int width,
+                                const int num_bbox,
+                                const int channels_trans,
+                                const int no_trans,
+                                const float spatial_scale,
+                                const int output_dim,
+                                const int group_size,
+                                const int pooled_size,
+                                const int part_size,
+                                const int sample_per_part,
+                                const float trans_std);
+
+void DeformablePSROIPoolBackwardAcc(const at::Tensor out_grad,
+                                    const at::Tensor data,
+                                    const at::Tensor bbox,
+                                    const at::Tensor trans,
+                                    const at::Tensor top_count,
+                                    at::Tensor in_grad,
+                                    at::Tensor trans_grad,
+                                    const int batch,
+                                    const int channels,
+                                    const int height,
+                                    const int width,
+                                    const int num_bbox,
+                                    const int channels_trans,
+                                    const int no_trans,
+                                    const float spatial_scale,
+                                    const int output_dim,
+                                    const int group_size,
+                                    const int pooled_size,
+                                    const int part_size,
+                                    const int sample_per_part,
+                                    const float trans_std);
+
+
+void modulated_deform_conv_cuda_forward(at::Tensor input, at::Tensor weight,
+                         at::Tensor bias, at::Tensor ones,
+                         at::Tensor offset, at::Tensor mask,
+                         at::Tensor output, at::Tensor columns,
+                         int kernel_h, int kernel_w,
+                         const int stride_h, const int stride_w,
+                         const int pad_h, const int pad_w,
+                         const int dilation_h, const int dilation_w,
+                         const int deformable_group)
+{
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 8, input, weight, bias, ones, offset, mask, output, columns));
+    AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+    AT_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
+
+    const int batch = input.size(0);
+    const int channels = input.size(1);
+    const int height = input.size(2);
+    const int width = input.size(3);
+
+    const int channels_out = weight.size(0);
+    const int channels_kernel = weight.size(1);
+    const int kernel_h_ = weight.size(2);
+    const int kernel_w_ = weight.size(3);
+
+    if (kernel_h_ != kernel_h || kernel_w_ != kernel_w)
+        AT_ERROR("Input shape and kernel shape wont match: (%d x %d vs %d x %d).",
+        kernel_h_, kernel_w, kernel_h_, kernel_w_);
+    if (channels != channels_kernel)
+        AT_ERROR("Input shape and kernel channels wont match: (%d vs %d).",
+        channels, channels_kernel);
+
+    const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
+    const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+    if (ones.ndimension() != 2 ||
+        ones.size(0) * ones.size(1) < height_out * width_out)
+    {
+        // Resize plane and fill with ones...
+        ones = at::ones({height_out, width_out}, input.type());
+    }
+
+    // resize output
+    output = output.view({batch, channels_out, height_out, width_out});
+    // resize temporary columns
+    columns = at::zeros({channels * kernel_h * kernel_w, 1 * height_out * width_out}, input.type());
+
+    for (int b = 0; b < batch; b++)
+    {
+        // Do Bias first:
+        // M,N,K are dims of matrix A and B
+        // (see http://docs.nvidia.com/cuda/cublas/#cublas-lt-t-gt-gemm)
+        // (N x 1) (1 x M)
+        output[b] = output[b].flatten(1).addmm_(bias.view({-1, 1}), ones.view({1, -1}), 0.0f, 1.0f).view_as(output[b]);
+
+        modulated_deformable_im2col_cuda(input[b], offset[b], mask[b],
+                                         1, channels, height, width,
+                                         height_out, width_out, kernel_h, kernel_w,
+                                         pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,
+                                         deformable_group, columns);
+
+        //(k * m)  x  (m * n)
+        // Y = WC
+        output[b] = output[b].flatten(1).addmm_(weight.flatten(1), columns).view_as(output[b]);
+    }
+}
+
+void modulated_deform_conv_cuda_backward(at::Tensor input, at::Tensor weight,
+                          at::Tensor bias, at::Tensor ones,
+                          at::Tensor offset, at::Tensor mask,
+                          at::Tensor columns,
+                          at::Tensor grad_input, at::Tensor grad_weight,
+                          at::Tensor grad_bias, at::Tensor grad_offset,
+                          at::Tensor grad_mask, at::Tensor grad_output,
+                          int kernel_h, int kernel_w,
+                          int stride_h, int stride_w,
+                          int pad_h, int pad_w,
+                          int dilation_h, int dilation_w,
+                          int deformable_group)
+{
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 13, input, weight, bias, ones, offset, mask, columns,
+    //                                        grad_input, grad_weight, grad_bias, grad_offset, grad_mask, grad_output));
+    AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+    AT_CHECK(weight.is_contiguous(), "weight tensor has to be contiguous");
+
+    const int batch = input.size(0);
+    const int channels = input.size(1);
+    const int height = input.size(2);
+    const int width = input.size(3);
+
+    const int channels_kernel = weight.size(1);
+    const int kernel_h_ = weight.size(2);
+    const int kernel_w_ = weight.size(3);
+    if (kernel_h_ != kernel_h || kernel_w_ != kernel_w)
+        AT_ERROR("Input shape and kernel shape wont match: (%d x %d vs %d x %d).",
+        kernel_h_, kernel_w, kernel_h_, kernel_w_);
+    if (channels != channels_kernel)
+        AT_ERROR("Input shape and kernel channels wont match: (%d vs %d).",
+        channels, channels_kernel);
+
+    const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;
+    const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;
+
+    if (ones.ndimension() != 2 ||
+        ones.size(0) * ones.size(1) < height_out * width_out)
+    {
+        // Resize plane and fill with ones...
+        ones = at::ones({height_out, width_out}, input.type());
+    }
+
+    grad_input = grad_input.view({batch, channels, height, width});
+    columns = at::zeros({channels * kernel_h * kernel_w, height_out * width_out}, input.type());
+
+    for (int b = 0; b < batch; b++)
+    {
+        columns.addmm_(weight.flatten(1).transpose(0, 1), grad_output[b].flatten(1), 0.0f, 1.0f);
+
+        // gradient w.r.t. input coordinate data
+        modulated_deformable_col2im_coord_cuda(columns, input[b], offset[b], mask[b],
+                                               1, channels, height, width,
+                                               height_out, width_out, kernel_h, kernel_w,
+                                               pad_h, pad_w, stride_h, stride_w,
+                                               dilation_h, dilation_w, deformable_group,
+                                               grad_offset[b], grad_mask[b]);
+        // gradient w.r.t. input data
+        modulated_deformable_col2im_cuda(columns, offset[b], mask[b],
+                                         1, channels, height, width,
+                                         height_out, width_out, kernel_h, kernel_w,
+                                         pad_h, pad_w, stride_h, stride_w,
+                                         dilation_h, dilation_w, deformable_group,
+                                         grad_input[b]);
+
+        // gradient w.r.t. weight, dWeight should accumulate across the batch and group
+        modulated_deformable_im2col_cuda(input[b], offset[b], mask[b],
+                                         1, channels, height, width,
+                                         height_out, width_out, kernel_h, kernel_w,
+                                         pad_h, pad_w, stride_h, stride_w,
+                                         dilation_h, dilation_w, deformable_group,
+                                         columns);
+
+        grad_weight = grad_weight.flatten(1).addmm_(grad_output[b].flatten(1), columns.transpose(0, 1)).view_as(grad_weight);
+
+        grad_bias = grad_bias.view({-1, 1}).addmm_(grad_output[b].flatten(1), ones.view({-1, 1})).view(-1);
+    }
+
+}
+
+void deform_psroi_pooling_cuda_forward(at::Tensor  input, at::Tensor  bbox,
+                                       at::Tensor  trans,
+                                       at::Tensor  out, at::Tensor  top_count,
+                                       const int no_trans,
+                                       const float spatial_scale,
+                                       const int output_dim,
+                                       const int group_size,
+                                       const int pooled_size,
+                                       const int part_size,
+                                       const int sample_per_part,
+                                       const float trans_std)
+{
+    AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 5, input, bbox, trans, out, top_count));
+
+    const int batch = input.size(0);
+    const int channels = input.size(1);
+    const int height = input.size(2);
+    const int width = input.size(3);
+    const int channels_trans = no_trans? 2 : trans.size(1);
+
+    const int num_bbox = bbox.size(0);
+    if (num_bbox != out.size(0))
+        AT_ERROR("Output shape and bbox number wont match: (%d vs %d).",
+                out.size(0), num_bbox);
+
+    DeformablePSROIPoolForward(input, bbox, trans, out, top_count,
+                               batch, channels, height, width,
+                               num_bbox,
+                               channels_trans,
+                               no_trans,
+                               spatial_scale,
+                               output_dim,
+                               group_size,
+                               pooled_size,
+                               part_size,
+                               sample_per_part,
+                               trans_std);
+}
+
+void deform_psroi_pooling_cuda_backward(at::Tensor  out_grad,
+                                        at::Tensor  input, at::Tensor  bbox,
+                                        at::Tensor  trans, at::Tensor  top_count,
+                                        at::Tensor  input_grad, at::Tensor  trans_grad,
+                                        const int no_trans,
+                                        const float spatial_scale,
+                                        const int output_dim,
+                                        const int group_size,
+                                        const int pooled_size,
+                                        const int part_size,
+                                        const int sample_per_part,
+                                        const float trans_std)
+{
+    AT_CHECK(out_grad.is_contiguous(), "out_grad tensor has to be contiguous");
+    AT_CHECK(input.is_contiguous(), "input tensor has to be contiguous");
+    // THCAssertSameGPU(THCudaTensor_checkGPU(state, 7, input, bbox, trans, out_grad, top_count,
+    //                 input_grad, trans_grad));
+
+    const int batch = input.size(0);
+    const int channels = input.size(1);
+    const int height = input.size(2);
+    const int width = input.size(3);
+    const int channels_trans = no_trans? 2 : trans.size(1);
+
+    const int num_bbox = bbox.size(0);
+    if (num_bbox != out_grad.size(0))
+        AT_ERROR("Output shape and bbox number wont match: (%d vs %d).",
+                out_grad.size(0), num_bbox);
+
+    DeformablePSROIPoolBackwardAcc(out_grad,
+                                   input,
+                                   bbox,
+                                   trans,
+                                   top_count,
+                                   input_grad,
+                                   trans_grad,
+                                   batch, channels, height, width, num_bbox,
+                                   channels_trans,
+                                   no_trans,
+                                   spatial_scale,
+                                   output_dim,
+                                   group_size,
+                                   pooled_size,
+                                   part_size,
+                                   sample_per_part,
+                                   trans_std);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("modulated_deform_conv_cuda_forward", &modulated_deform_conv_cuda_forward,
+          "modulated deform conv forward (CUDA)");
+    m.def("modulated_deform_conv_cuda_backward", &modulated_deform_conv_cuda_backward,
+          "modulated deform conv backward (CUDA)");
+    m.def("deform_psroi_pooling_cuda_forward", &deform_psroi_pooling_cuda_forward,
+          "deform psroi pooling forward(CUDA)");
+    m.def("deform_psroi_pooling_cuda_backward", &deform_psroi_pooling_cuda_backward,
+          "deform psroi pooling backward(CUDA)");
+}
\ No newline at end of file