diff --git a/mmdet/ops/affine_grid/__init__.py b/mmdet/ops/affine_grid/__init__.py new file mode 100644 index 00000000000..8530ade3384 --- /dev/null +++ b/mmdet/ops/affine_grid/__init__.py @@ -0,0 +1,3 @@ +from .affine_grid import affine_grid + +__all__ = ['affine_grid'] diff --git a/mmdet/ops/affine_grid/affine_grid.py b/mmdet/ops/affine_grid/affine_grid.py new file mode 100644 index 00000000000..94bacb5e596 --- /dev/null +++ b/mmdet/ops/affine_grid/affine_grid.py @@ -0,0 +1,68 @@ +import torch +import torch.nn.functional as F +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +from . import affine_grid_cuda + + +class _AffineGridGenerator(Function): + + @staticmethod + def forward(ctx, theta, size, align_corners): + + ctx.save_for_backward(theta) + ctx.size = size + ctx.align_corners = align_corners + + func = affine_grid_cuda.affine_grid_generator_forward + + output = func(theta, size, align_corners) + + return output + + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + theta = ctx.saved_tensors + size = ctx.size + align_corners = ctx.align_corners + + func = affine_grid_cuda.affine_grid_generator_backward + + grad_input = func(grad_output, theta, size, align_corners) + + return grad_input, None, None + + +def affine_grid(theta, size, align_corners=False): + if torch.__version__ >= '1.3': + return F.affine_grid(theta, size, align_corners) + elif align_corners: + return F.affine_grid(theta, size) + else: + # enforce floating point dtype on theta + if not theta.is_floating_point(): + raise ValueError( + 'Expected theta to have floating point type, but got {}'. + format(theta.dtype)) + # check that shapes and sizes match + if len(size) == 4: + if theta.dim() != 3 or theta.size(-2) != 2 or theta.size(-1) != 3: + raise ValueError( + 'Expected a batch of 2D affine matrices of shape Nx2x3 ' + 'for size {}. Got {}.'.format(size, theta.shape)) + elif len(size) == 5: + if theta.dim() != 3 or theta.size(-2) != 3 or theta.size(-1) != 4: + raise ValueError( + 'Expected a batch of 3D affine matrices of shape Nx3x4 ' + 'for size {}. Got {}.'.format(size, theta.shape)) + else: + raise NotImplementedError( + 'affine_grid only supports 4D and 5D sizes, ' + 'for 2D and 3D affine transforms, respectively. ' + 'Got size {}.'.format(size)) + if min(size) <= 0: + raise ValueError( + 'Expected non-zero, positive output size. Got {}'.format(size)) + return _AffineGridGenerator.apply(theta, size, align_corners) diff --git a/mmdet/ops/affine_grid/src/affine_grid_cuda.cpp b/mmdet/ops/affine_grid/src/affine_grid_cuda.cpp new file mode 100644 index 00000000000..3874128c537 --- /dev/null +++ b/mmdet/ops/affine_grid/src/affine_grid_cuda.cpp @@ -0,0 +1,115 @@ +// Modified from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/AffineGridGenerator.cpp +#include +#include +#include + +namespace mmdetection { + +using namespace at; + +at::Tensor linspace_from_neg_one(const Tensor& grid, int64_t num_steps, + bool align_corners) { + if (num_steps <= 1) { + return at::tensor(0, grid.options()); + } + auto range = at::linspace(-1, 1, num_steps, grid.options()); + if (!align_corners) { + range = range * (num_steps - 1) / num_steps; + } + return range; +} + +Tensor make_base_grid_4D(const Tensor& theta, int64_t N, int64_t C, int64_t H, + int64_t W, bool align_corners) { + auto base_grid = at::empty({N, H, W, 3}, theta.options()); + + base_grid.select(-1, 0).copy_(linspace_from_neg_one(theta, W, align_corners)); + base_grid.select(-1, 1).copy_( + linspace_from_neg_one(theta, H, align_corners).unsqueeze_(-1)); + base_grid.select(-1, 2).fill_(1); + + return base_grid; +} + +Tensor make_base_grid_5D(const Tensor& theta, int64_t N, int64_t C, int64_t D, + int64_t H, int64_t W, bool align_corners) { + auto base_grid = at::empty({N, D, H, W, 4}, theta.options()); + + base_grid.select(-1, 0).copy_(linspace_from_neg_one(theta, W, align_corners)); + base_grid.select(-1, 1).copy_( + linspace_from_neg_one(theta, H, align_corners).unsqueeze_(-1)); + base_grid.select(-1, 2).copy_(linspace_from_neg_one(theta, D, align_corners) + .unsqueeze_(-1) + .unsqueeze_(-1)); + base_grid.select(-1, 3).fill_(1); + + return base_grid; +} + +Tensor affine_grid_generator_4D_forward(const Tensor& theta, int64_t N, + int64_t C, int64_t H, int64_t W, + bool align_corners) { + Tensor base_grid = make_base_grid_4D(theta, N, C, H, W, align_corners); + auto grid = base_grid.view({N, H * W, 3}).bmm(theta.transpose(1, 2)); + return grid.view({N, H, W, 2}); +} + +Tensor affine_grid_generator_5D_forward(const Tensor& theta, int64_t N, + int64_t C, int64_t D, int64_t H, + int64_t W, bool align_corners) { + Tensor base_grid = make_base_grid_5D(theta, N, C, D, H, W, align_corners); + auto grid = base_grid.view({N, D * H * W, 4}).bmm(theta.transpose(1, 2)); + return grid.view({N, D, H, W, 3}); +} + +Tensor affine_grid_generator_forward(const Tensor& theta, IntArrayRef size, + bool align_corners) { + if (size.size() == 4) { + return affine_grid_generator_4D_forward(theta, size[0], size[1], size[2], + size[3], align_corners); + } else { + return affine_grid_generator_5D_forward(theta, size[0], size[1], size[2], + size[3], size[4], align_corners); + } +} + +Tensor affine_grid_generator_4D_backward(const Tensor& grad_grid, int64_t N, + int64_t C, int64_t H, int64_t W, + bool align_corners) { + auto base_grid = make_base_grid_4D(grad_grid, N, C, H, W, align_corners); + AT_ASSERT(grad_grid.sizes() == IntArrayRef({N, H, W, 2})); + auto grad_theta = base_grid.view({N, H * W, 3}) + .transpose(1, 2) + .bmm(grad_grid.view({N, H * W, 2})); + return grad_theta.transpose(1, 2); +} + +Tensor affine_grid_generator_5D_backward(const Tensor& grad_grid, int64_t N, + int64_t C, int64_t D, int64_t H, + int64_t W, bool align_corners) { + auto base_grid = make_base_grid_5D(grad_grid, N, C, D, H, W, align_corners); + AT_ASSERT(grad_grid.sizes() == IntArrayRef({N, D, H, W, 3})); + auto grad_theta = base_grid.view({N, D * H * W, 4}) + .transpose(1, 2) + .bmm(grad_grid.view({N, D * H * W, 3})); + return grad_theta.transpose(1, 2); +} + +Tensor affine_grid_generator_backward(const Tensor& grad, IntArrayRef size, + bool align_corners) { + if (size.size() == 4) { + return affine_grid_generator_4D_backward(grad, size[0], size[1], size[2], + size[3], align_corners); + } else { + return affine_grid_generator_5D_backward(grad, size[0], size[1], size[2], + size[3], size[4], align_corners); + } +} +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("affine_grid_generator_forward", &affine_grid_generator_forward, + "affine_grid_generator_forward"); + m.def("affine_grid_generator_backward", &affine_grid_generator_backward, + "affine_grid_generator_backward"); +} + +} // namespace mmdetection diff --git a/mmdet/ops/grid_sampler/grid_sampler.py b/mmdet/ops/grid_sampler/grid_sampler.py index d359d8dfe95..a112fa82719 100644 --- a/mmdet/ops/grid_sampler/grid_sampler.py +++ b/mmdet/ops/grid_sampler/grid_sampler.py @@ -61,8 +61,10 @@ def grid_sample(input, mode='bilinear', padding_mode='zeros', align_corners=False): - if torch.__version__ >= '1.3' or align_corners: + if torch.__version__ >= '1.3': return F.grid_sample(input, grid, mode, padding_mode, align_corners) + elif align_corners: + return F.grid_sample(input, grid, mode, padding_mode) else: # use self-compiled grid_sampler to support align_corners=False diff --git a/setup.py b/setup.py index 9704b4d614c..a11263f22cc 100755 --- a/setup.py +++ b/setup.py @@ -269,6 +269,10 @@ def gen_packages_items(): sources=[ 'src/masked_conv2d_cuda.cpp', 'src/masked_conv2d_kernel.cu' ]), + make_cuda_ext( + name='affine_grid_cuda', + module='mmdet.ops.affine_grid', + sources=['src/affine_grid_cuda.cpp']), make_cuda_ext( name='grid_sampler_cuda', module='mmdet.ops.grid_sampler',