Update DeformConv to be more consistent w/ Conv2d

* rename some variables and arguments to match Conv2d;
* add optional bias;
* add weight, offset and bias as module parameters;
* remove the n_parallel_imgs parameter;
* Fix __repr__;
* etc..

Initialization of weight and bias is the same as in Conv2d, and
initialization of offsets to zero is the same as in the paper.

This also includes some other small unrelated fixes/improvements.

Author: Pedro Freire

test/test_ops.py

@@ -1,14 +1,14 @@
 from __future__ import division
 import math
-from typing import Tuple
 import unittest
 
 import numpy as np
 
 import torch
+from torch import Tensor
 from torch.autograd import gradcheck
+from torch.jit.annotations import Tuple
 from torch.nn.modules.utils import _pair
-from torch import Tensor
 from torchvision import ops
 
 
@@ -374,44 +374,45 @@ def test_nms_cuda(self):
 
 
 class DeformConvTester(OpTester, unittest.TestCase):
-    def expected_fn(self, x, offsets, weights, *args, stride=1, pad=0, dilation=1):
+    def expected_fn(self, x, weight, offset, bias, stride=1, padding=0, dilation=1):
         stride_h, stride_w = _pair(stride)
-        pad_h, pad_w = _pair(pad)
+        pad_h, pad_w = _pair(padding)
         dil_h, dil_w = _pair(dilation)
-        weights_h, weights_w = weights.shape[-2:]
+        weight_h, weight_w = weight.shape[-2:]
 
         n_batches, n_in_channels, in_h, in_w = x.shape
-        n_out_channels = weights.shape[0]
+        n_out_channels = weight.shape[0]
 
-        out_h = (in_h + 2 * pad_h - (dil_h * (weights_h - 1) + 1)) // stride_h + 1
-        out_w = (in_w + 2 * pad_w - (dil_w * (weights_w - 1) + 1)) // stride_w + 1
+        out_h = (in_h + 2 * pad_h - (dil_h * (weight_h - 1) + 1)) // stride_h + 1
+        out_w = (in_w + 2 * pad_w - (dil_w * (weight_w - 1) + 1)) // stride_w + 1
 
-        n_offset_grps = offsets.shape[1] // (2 * weights_h * weights_w)
+        n_offset_grps = offset.shape[1] // (2 * weight_h * weight_w)
         in_c_per_offset_grp = n_in_channels // n_offset_grps
 
-        n_weight_grps = n_in_channels // weights.shape[1]
-        in_c_per_weight_grp = weights.shape[1]
+        n_weight_grps = n_in_channels // weight.shape[1]
+        in_c_per_weight_grp = weight.shape[1]
         out_c_per_weight_grp = n_out_channels // n_weight_grps
 
         out = torch.zeros(n_batches, n_out_channels, out_h, out_w, device=x.device, dtype=x.dtype)
         for b in range(n_batches):
             for c_out in range(n_out_channels):
                 for i in range(out_h):
                     for j in range(out_w):
-                        for di in range(weights_h):
-                            for dj in range(weights_w):
+                        for di in range(weight_h):
+                            for dj in range(weight_w):
                                 for c in range(in_c_per_weight_grp):
                                     weight_grp = c_out // out_c_per_weight_grp
                                     c_in = weight_grp * in_c_per_weight_grp + c
 
                                     offset_grp = c_in // in_c_per_offset_grp
-                                    offset_idx = 2 * (offset_grp * (weights_h * weights_w) + di * weights_w + dj)
+                                    offset_idx = 2 * (offset_grp * (weight_h * weight_w) + di * weight_w + dj)
 
-                                    pi = stride_h * i - pad_h + dil_h * di + offsets[b, offset_idx, i, j]
-                                    pj = stride_w * j - pad_w + dil_w * dj + offsets[b, offset_idx + 1, i, j]
+                                    pi = stride_h * i - pad_h + dil_h * di + offset[b, offset_idx, i, j]
+                                    pj = stride_w * j - pad_w + dil_w * dj + offset[b, offset_idx + 1, i, j]
 
-                                    out[b, c_out, i, j] += (weights[c_out, c, di, dj] *
+                                    out[b, c_out, i, j] += (weight[c_out, c, di, dj] *
                                                             bilinear_interpolate(x[b, c_in, :, :], pi, pj))
+        out += bias.view(1, n_out_channels, 1, 1)
         return out
 
     def get_fn_args(self, device, contiguous):
@@ -442,36 +443,50 @@ def get_fn_args(self, device, contiguous):
         weight = torch.randn(n_out_channels, n_in_channels // n_weight_grps, weight_h, weight_w,
                              device=device, dtype=self.dtype, requires_grad=True)
 
+        bias = torch.randn(n_out_channels, device=device, dtype=self.dtype, requires_grad=True)
+
         if not contiguous:
             x = x.permute(0, 1, 3, 2).contiguous().permute(0, 1, 3, 2)
             offset = offset.permute(1, 3, 0, 2).contiguous().permute(2, 0, 3, 1)
             weight = weight.permute(3, 2, 0, 1).contiguous().permute(2, 3, 1, 0)
 
-        return x, offset, weight, stride, pad, dilation
+        return x, weight, offset, bias, stride, pad, dilation
 
     def _test_forward(self, device, contiguous):
-        x, offset, weight, stride, pad, dilation = self.get_fn_args(device, contiguous)
+        x, _, _, _, stride, padding, dilation = self.get_fn_args(device, contiguous)
+        in_channels = 6
+        out_channels = 2
+        kernel_size = (3, 2)
+        groups = 2
+        offset_groups = 3
+
+        layer = ops.DeformConv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
+                                 dilation=dilation, groups=groups, offset_groups=offset_groups)
+        layer.offset_conv.weight.data = torch.randn_like(layer.offset_conv.weight.data)
+        res = layer(x)
 
-        res = ops.DeformConv(stride=stride, pad=pad, dilation=dilation)(x, offset, weight)
-        expected = self.expected_fn(x, offset, weight, stride=stride, pad=pad, dilation=dilation)
+        weight = layer.weight.data.to(device=x.device, dtype=x.dtype)
+        offset = layer.offset_conv.to(device=x.device, dtype=x.dtype)(x)
+        bias = layer.bias.data.to(device=x.device, dtype=x.dtype)
+        expected = self.expected_fn(x, weight, offset, bias, stride=stride, padding=padding, dilation=dilation)
 
-        self.assertTrue(torch.allclose(res, expected), '\nres:\n{}\nexpected:\n{}'.format(x, res, expected))
+        self.assertTrue(torch.allclose(res, expected), '\nres:\n{}\nexpected:\n{}'.format(res, expected))
 
     def _test_backward(self, device, contiguous):
-        x, offset, weight, stride, pad, dilation = self.get_fn_args(device, contiguous)
+        x, weight, offset, bias, stride, padding, dilation = self.get_fn_args(device, contiguous)
 
-        def func(x_, offset_, weight_):
-            return ops.deform_conv(x_, offset_, weight_, stride=stride, pad=pad, dilation=dilation)
+        def func(x_, weight_, offset_, bias_):
+            return ops.deform_conv2d(x_, weight_, offset_, bias_, stride=stride, padding=padding, dilation=dilation)
 
-        gradcheck(func, (x, offset, weight), nondet_tol=1e-5)
+        gradcheck(func, (x, weight, offset, bias), nondet_tol=1e-5)
 
         @torch.jit.script
-        def script_func(x_, offset_, weight_, stride_, pad_, dilation_):
-            # type: (Tensor, Tensor, Tensor, Tuple[int, int], Tuple[int, int], Tuple[int, int]) -> Tensor
-            return ops.deform_conv(x_, offset_, weight_, stride=stride_, pad=pad_, dilation=dilation_)
+        def script_func(x_, weight_, offset_, bias_, stride_, pad_, dilation_):
+            # type: (Tensor, Tensor, Tensor, Tensor, Tuple[int, int], Tuple[int, int], Tuple[int, int]) -> Tensor
+            return ops.deform_conv2d(x_, weight_, offset_, bias_, stride=stride_, padding=pad_, dilation=dilation_)
 
-        gradcheck(lambda z, off, wei: script_func(z, off, wei, stride, pad, dilation),
-                  (x, offset, weight), nondet_tol=1e-5)
+        gradcheck(lambda z, wei, off, bi: script_func(z, wei, off, bi, stride, padding, dilation),
+                  (x, weight, offset, bias), nondet_tol=1e-5)
 
 
 if __name__ == '__main__':

torchvision/csrc/DeformConv.h

@@ -6,49 +6,48 @@
 #include "cuda/vision_cuda.h"
 #endif
 
-at::Tensor DCN_forward(
+at::Tensor DeformConv2d_forward(
     const Tensor& input,
+    const Tensor& weight,
     const Tensor& offset,
-    const Tensor& weights,
+    const Tensor& bias,
     const std::pair<int, int>& stride,
-    const std::pair<int, int>& pad,
+    const std::pair<int, int>& padding,
     const std::pair<int, int>& dilation,
-    const int groups,
-    const int deformable_groups,
-    const int n_parallel_imgs) {
+    const int groups, const int offset_groups) {
   if (input.type().is_cuda()) {
 #ifdef WITH_CUDA
-    return DCN_forward_cuda(input.contiguous(), offset.contiguous(), weights.contiguous(), stride, pad,
-                      dilation, groups, deformable_groups, n_parallel_imgs);
+    return DeformConv2d_forward_cuda(input.contiguous(), weight.contiguous(), offset.contiguous(),
+            bias.contiguous(), stride, padding, dilation, groups, offset_groups);
 #else
     AT_ERROR("Not compiled with GPU support");
 #endif
   }
-  return DCN_forward_cpu(input.contiguous(), offset.contiguous(), weights.contiguous(), stride, pad,
-                    dilation, groups, deformable_groups, n_parallel_imgs);
+  return DeformConv2d_forward_cpu(input.contiguous(), weight.contiguous(), offset.contiguous(),
+            bias.contiguous(), stride, padding, dilation, groups, offset_groups);
 }
 
-std::tuple<at::Tensor, at::Tensor, at::Tensor> DCN_backward(
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> DeformConv2d_backward(
     const at::Tensor& grad,
     const Tensor& input,
+    const Tensor& weight,
     const Tensor& offset,
-    const Tensor& weights,
+    const Tensor& bias,
     const std::pair<int, int>& stride,
-    const std::pair<int, int>& pad,
+    const std::pair<int, int>& padding,
     const std::pair<int, int>& dilation,
     const int groups,
-    const int deformable_groups,
-    const int n_parallel_imgs) {
+    const int offset_groups) {
   if (grad.type().is_cuda()) {
 #ifdef WITH_CUDA
-    return DCN_backward_cuda(grad.contiguous(), input.contiguous(), offset.contiguous(), weights.contiguous(), stride, pad,
-                      dilation, groups, deformable_groups, n_parallel_imgs);
+    return DeformConv2d_backward_cuda(grad.contiguous(), input.contiguous(), weight.contiguous(), offset.contiguous(),
+            bias.contiguous(), stride, padding, dilation, groups, offset_groups);
 #else
     AT_ERROR("Not compiled with GPU support");
 #endif
   }
-  return DCN_backward_cpu(grad.contiguous(), input.contiguous(), offset.contiguous(), weights.contiguous(), stride, pad,
-                      dilation, groups, deformable_groups, n_parallel_imgs);
+  return DeformConv2d_backward_cpu(grad.contiguous(), input.contiguous(), weight.contiguous(), offset.contiguous(),
+          bias.contiguous(), stride, padding, dilation, groups, offset_groups);
 }
 
 using namespace at;
@@ -57,35 +56,35 @@ using torch::autograd::AutogradContext;
 using torch::autograd::Variable;
 using torch::autograd::variable_list;
 
-class DeformConvFunction : public torch::autograd::Function<DeformConvFunction> {
+class DeformConv2dFunction : public torch::autograd::Function<DeformConv2dFunction> {
  public:
   static variable_list forward(
       AutogradContext* ctx,
       Variable input,
+      Variable weight,
       Variable offset,
-      Variable weights,
+      Variable bias,
       int64_t stride_h, int64_t stride_w,
       int64_t pad_h, int64_t pad_w,
       int64_t dilation_h, int64_t dilation_w,
       int64_t groups,
-      int64_t deformable_groups,
-      int64_t n_parallel_imgs) {
-    auto output = DCN_forward(input, offset, weights,
+      int64_t offset_groups) {
+    auto output = DeformConv2d_forward(
+        input, weight, offset, bias,
         {stride_h, stride_w},
         {pad_h, pad_w},
         {dilation_h, dilation_w},
-        groups, deformable_groups, n_parallel_imgs);
+        groups, offset_groups);
 
-    ctx->save_for_backward({input, offset, weights});
+    ctx->save_for_backward({input, weight, offset, bias});
     ctx->saved_data["stride_h"] = stride_h;
     ctx->saved_data["stride_w"] = stride_w;
     ctx->saved_data["pad_h"] = pad_h;
     ctx->saved_data["pad_w"] = pad_w;
     ctx->saved_data["dilation_h"] = dilation_h;
     ctx->saved_data["dilation_w"] = dilation_w;
     ctx->saved_data["groups"] = groups;
-    ctx->saved_data["deformable_groups"] = deformable_groups;
-    ctx->saved_data["n_parallel_imgs"] = n_parallel_imgs;
+    ctx->saved_data["offset_groups"] = offset_groups;
 
     return {output,};
   }
@@ -95,8 +94,9 @@ class DeformConvFunction : public torch::autograd::Function<DeformConvFunction>
       variable_list grad_output) {
     auto saved = ctx->get_saved_variables();
     auto input = saved[0];
-    auto offset = saved[1];
-    auto weight = saved[2];
+    auto weight = saved[1];
+    auto offset = saved[2];
+    auto bias = saved[3];
 
     auto stride_h = ctx->saved_data["stride_h"].toInt();
     auto stride_w = ctx->saved_data["stride_w"].toInt();
@@ -105,37 +105,36 @@ class DeformConvFunction : public torch::autograd::Function<DeformConvFunction>
     auto dilation_h = ctx->saved_data["dilation_h"].toInt();
     auto dilation_w = ctx->saved_data["dilation_w"].toInt();
     auto groups = ctx->saved_data["groups"].toInt();
-    auto deformable_groups = ctx->saved_data["deformable_groups"].toInt();
-    auto n_parallel_imgs = ctx->saved_data["n_parallel_imgs"].toInt();
+    auto offset_groups = ctx->saved_data["offset_groups"].toInt();
 
-    auto grads = DCN_backward(grad_output[0],
-        input, offset, weight,
+    auto grads = DeformConv2d_backward(grad_output[0],
+        input, weight, offset, bias,
         {stride_h, stride_w},
         {pad_h, pad_w},
         {dilation_h, dilation_w},
-        groups, deformable_groups, n_parallel_imgs);
+        groups, offset_groups);
     auto grad_input = std::get<0>(grads);
-    auto grad_offset = std::get<1>(grads);
-    auto grad_weight = std::get<2>(grads);
+    auto grad_weight = std::get<1>(grads);
+    auto grad_offset = std::get<2>(grads);
+    auto grad_bias = std::get<3>(grads);
 
-    return {grad_input, grad_offset, grad_weight,
-            Variable(), Variable(), Variable(),
+    return {grad_input, grad_weight, grad_offset,
+            grad_bias, Variable(), Variable(),
             Variable(), Variable(), Variable(),
             Variable(), Variable(), Variable(),};
   }
 };
 
-Tensor deform_conv(
+Tensor deform_conv2d(
     const Tensor& input,
+    const Tensor& weight,
     const Tensor& offset,
-    const Tensor& weights,
+    const Tensor& bias,
     int64_t stride_h, int64_t stride_w,
     int64_t pad_h, int64_t pad_w,
     int64_t dilation_h, int64_t dilation_w,
-    int64_t groups,
-    int64_t deformable_groups,
-    int64_t n_parallel_imgs) {
-  auto result = DeformConvFunction::apply(input, offset, weights, stride_h, stride_w, pad_h, pad_w,
-                          dilation_h, dilation_w, groups, deformable_groups, n_parallel_imgs);
+    int64_t groups, int64_t offset_groups) {
+  auto result = DeformConv2dFunction::apply(input, weight, offset, bias, stride_h, stride_w, pad_h, pad_w,
+                                            dilation_h, dilation_w, groups, offset_groups);
   return result[0];
 }