tvdcn 0.4.0

Torchvision+ Deformable Convolutional Networks

Torchvision+ Deformable Convolution Networks

This package contains the PyTorch implementations of the Deformable Convolution operation (the commonly used torchvision.ops.deform_conv2d) proposed in https://arxiv.org/abs/1811.11168, and the Transposed Deformable Convolution proposed in https://arxiv.org/abs/2210.09446 (currently without interpolation kernel scaling). It also supports their 1D and 3D equivalences, which are not available in torchvision (thus the name).

Highlights

• Supported operations: (All are implemented in C++/Cuda)

  • tvdcn.ops.deform_conv1d
  • tvdcn.ops.deform_conv2d
  • tvdcn.ops.deform_conv3d
  • tvdcn.ops.deform_conv_transpose1d
  • tvdcn.ops.deform_conv_transpose2d
  • tvdcn.ops.deform_conv_transpose3d

• And the following supplementary operations (mask activation proposed in https://arxiv.org/abs/2211.05778):

  • tvdcn.ops.mask_softmax1d
  • tvdcn.ops.mask_softmax2d
  • tvdcn.ops.mask_softmax3d

• Both offset and mask can be turned off, and can be applied in separate groups.

• All the nn.Module wrappers for these operations are implemented, everything is @torch.jit.script-able! Please check Usage.

Note: We don't care much about onnx exportation, but if you do, you can check this repo: https://github.com/masamitsu-murase/deform_conv2d_onnx_exporter.

Requirements

• torch>=1.9.0

Installation

From PyPI:

tvdcn provides some prebuilt wheels on PyPI. Run this command to install:

pip install tvdcn

The Linux and Windows wheels are built with Cuda 11.8. If you cannot find a wheel for your Arch/Python/Cuda, or there is any problem with library linking when importing, please proceed to instructions to build from source, all steps are super easy.

	Linux/Windows	MacOS
Python version:	3.8-3.11	3.8-3.11
PyTorch version:	torch==2.0.1	torch==2.0.1
Cuda version:	11.8	-
GPU CCs:	3.7,5.0,6.0,6.1,7.0,7.5,8.0,8.6,8.9,9.0+PTX	-

From Source:

For installing from source, you need a C++ compiler (gcc/msvc) and a Cuda compiler (nvcc) with C++17 features enabled. Clone this repo and execute the following command:

pip install .

Or just compile the binary for inplace usage:

python setup.py build_ext --inplace

A binary (.so file for Unix and .pyd file for Windows) should be compiled inside the tvdcn folder. To check if installation is successful, try:

import tvdcn

print('Library loaded successfully:', tvdcn.has_ops())
print('Compiled with Cuda:', tvdcn.with_cuda())

Note: We use soft Cuda version compatibility checking between the built binary and the installed PyTorch, which means only major version matching is required. However, we suggest building the binaries with the same Cuda version with installed PyTorch's Cuda version to prevent any possible conflict.

Usage

Functions:

Functionally, the package offers 6 functions (listed in Highlights) much similar to torchvision.ops.deform_conv2d. However, the order of parameters is slightly different, so be cautious (check this comparison). Specifically, the signatures of deform_conv2d and deform_conv_transpose2d look like this:

def deform_conv2d(
        input: Tensor,
        weight: Tensor,
        offset: Optional[Tensor] = None,
        mask: Optional[Tensor] = None,
        bias: Optional[Tensor] = None,
        stride: Union[int, Tuple[int, int]] = 1,
        padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
        groups: int = 1) -> Tensor:
    ...


def deform_conv_transpose2d(
        input: Tensor,
        weight: Tensor,
        offset: Optional[Tensor] = None,
        mask: Optional[Tensor] = None,
        bias: Optional[Tensor] = None,
        stride: Union[int, Tuple[int, int]] = 1,
        padding: Union[int, Tuple[int, int]] = 0,
        output_padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
        groups: int = 1) -> Tensor:
    ...

If offset=None and mask=None, the executed operations are identical to conventional convolution.

Neural Network Layers:

The nn.Module wrappers are:

• tvdcn.ops.DeformConv1d
• tvdcn.ops.DeformConv2d
• tvdcn.ops.DeformConv3d
• tvdcn.ops.DeformConvTranspose1d
• tvdcn.ops.DeformConvTranspose2d
• tvdcn.ops.DeformConvTranspose3d

They are subclasses of the torch.nn.modules._ConvNd, but you have to specify offset and optionally mask as extra inputs for the forward function. For example:

import torch

from tvdcn import DeformConv2d

input = torch.rand(2, 3, 64, 64)
offset = torch.rand(2, 2 * 3 * 3, 62, 62)
# if mask is None, perform the original deform_conv without modulation (v2)
mask = torch.rand(2, 1 * 3 * 3, 62, 62)

conv = DeformConv2d(3, 16, kernel_size=(3, 3))

output = conv(input, offset, mask)
print(output.shape)

Additionally, following many other implementations out there, we also implemented the packed wrappers:

• tvdcn.ops.PackedDeformConv1d
• tvdcn.ops.PackedDeformConv2d
• tvdcn.ops.PackedDeformConv3d
• tvdcn.ops.PackedDeformConvTranspose1d
• tvdcn.ops.PackedDeformConvTranspose2d
• tvdcn.ops.PackedDeformConvTranspose3d

These are easy-to-use classes that contain ordinary convolution layers with appropriate hyperparameters to generate offset (and mask if initialized with modulated=True); but that means less customization. The only tunable hyperparameters that effect these supplementary conv layers are offset_groups and mask_groups, which have been decoupled from and behave somewhat similar to groups.

To use the softmax activation for mask proposed in Deformable Convolution v3, set mask_activation='softmax'. offset_activation and mask_activation also accept any nn.Module.

import torch

from tvdcn import PackedDeformConv1d

input = torch.rand(2, 3, 128)

conv = PackedDeformConv1d(3, 16,
                          kernel_size=5,
                          modulated=True,
                          mask_activation='softmax')
# jit scripting
scripted_conv = torch.jit.script(conv)
print(scripted_conv)

output = scripted_conv(input)
print(output.shape)

Note: For transposed packed modules, we are generating offset and mask with pointwise convolution as we haven't found a better way to do it.

Check the examples folder, maybe you can find something helpful.

Acknowledgements

This for fun project is directly modified and extended from torchvision.ops.deform_conv2d.

License

The code is released under the MIT license. See LICENSE.txt for details.