Skip to content

Commit

Permalink
[WIP][New Op] Add deformable conv v2 (apache#16341)
Browse files Browse the repository at this point in the history
* add deformable conv v2

* fix lint and compiler warning

* fix lint

* fix pylint

* fix clang and lint

* fix base class, add test case

* fix gluon impl, add cpu forward

* address comment

* fix duplicate kernel name

* fix cpplint

* con't fix cpplint

* address comments

* fix mask scale

* make initial mask centered at 1 rather than 0.5

* fix submodule
  • Loading branch information
zhreshold authored and yajiedesign committed Nov 6, 2019
1 parent 40bc647 commit 76e21ff
Show file tree
Hide file tree
Showing 9 changed files with 1,839 additions and 2 deletions.
180 changes: 179 additions & 1 deletion python/mxnet/gluon/contrib/cnn/conv_layers.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
# pylint: disable= arguments-differ
"""Custom convolutional neural network layers in model_zoo."""

__all__ = ['DeformableConvolution']
__all__ = ['DeformableConvolution', 'ModulatedDeformableConvolution']

from .... import symbol
from ...block import HybridBlock
Expand Down Expand Up @@ -219,3 +219,181 @@ def __repr__(self):
return s.format(name=self.__class__.__name__,
mapping='{0} -> {1}'.format(shape[1] if shape[1] else None, shape[0]),
**self._kwargs_deformable_conv)


class ModulatedDeformableConvolution(HybridBlock):
"""2-D Deformable Convolution v2 (Dai, 2018).
The modulated deformable convolution operation is described in https://arxiv.org/abs/1811.11168
Parameters
----------
channels : int,
The dimensionality of the output space
i.e. the number of output channels in the convolution.
kernel_size : int or tuple/list of 2 ints, (Default value = (1,1))
Specifies the dimensions of the convolution window.
strides : int or tuple/list of 2 ints, (Default value = (1,1))
Specifies the strides of the convolution.
padding : int or tuple/list of 2 ints, (Default value = (0,0))
If padding is non-zero, then the input is implicitly zero-padded
on both sides for padding number of points.
dilation : int or tuple/list of 2 ints, (Default value = (1,1))
Specifies the dilation rate to use for dilated convolution.
groups : int, (Default value = 1)
Controls the connections between inputs and outputs.
At groups=1, all inputs are convolved to all outputs.
At groups=2, the operation becomes equivalent to having two convolution
layers side by side, each seeing half the input channels, and producing
half the output channels, and both subsequently concatenated.
num_deformable_group : int, (Default value = 1)
Number of deformable group partitions.
layout : str, (Default value = NCHW)
Dimension ordering of data and weight. Can be 'NCW', 'NWC', 'NCHW',
'NHWC', 'NCDHW', 'NDHWC', etc. 'N', 'C', 'H', 'W', 'D' stands for
batch, channel, height, width and depth dimensions respectively.
Convolution is performed over 'D', 'H', and 'W' dimensions.
use_bias : bool, (Default value = True)
Whether the layer for generating the output features uses a bias vector.
in_channels : int, (Default value = 0)
The number of input channels to this layer. If not specified,
initialization will be deferred to the first time `forward` is called
and input channels will be inferred from the shape of input data.
activation : str, (Default value = None)
Activation function to use. See :func:`~mxnet.ndarray.Activation`.
If you don't specify anything, no activation is applied
(ie. "linear" activation: `a(x) = x`).
weight_initializer : str or `Initializer`, (Default value = None)
Initializer for the `weight` weights matrix for the convolution layer
for generating the output features.
bias_initializer : str or `Initializer`, (Default value = zeros)
Initializer for the bias vector for the convolution layer
for generating the output features.
offset_weight_initializer : str or `Initializer`, (Default value = zeros)
Initializer for the `weight` weights matrix for the convolution layer
for generating the offset.
offset_bias_initializer : str or `Initializer`, (Default value = zeros),
Initializer for the bias vector for the convolution layer
for generating the offset.
offset_use_bias: bool, (Default value = True)
Whether the layer for generating the offset uses a bias vector.
Inputs:
- **data**: 4D input tensor with shape
`(batch_size, in_channels, height, width)` when `layout` is `NCHW`.
For other layouts shape is permuted accordingly.
Outputs:
- **out**: 4D output tensor with shape
`(batch_size, channels, out_height, out_width)` when `layout` is `NCHW`.
out_height and out_width are calculated as::
out_height = floor((height+2*padding[0]-dilation[0]*(kernel_size[0]-1)-1)/stride[0])+1
out_width = floor((width+2*padding[1]-dilation[1]*(kernel_size[1]-1)-1)/stride[1])+1
"""

def __init__(self, channels, kernel_size=(1, 1), strides=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1,
num_deformable_group=1, layout='NCHW', use_bias=True, in_channels=0, activation=None,
weight_initializer=None, bias_initializer='zeros',
offset_weight_initializer='zeros', offset_bias_initializer='zeros', offset_use_bias=True,
op_name='ModulatedDeformableConvolution', adj=None, prefix=None, params=None):
super(ModulatedDeformableConvolution, self).__init__(prefix=prefix, params=params)
with self.name_scope():
self._channels = channels
self._in_channels = in_channels

assert layout in ('NCHW', 'NHWC'), "Only supports 'NCHW' and 'NHWC' layout for now"
if isinstance(kernel_size, numeric_types):
kernel_size = (kernel_size,) * 2
if isinstance(strides, numeric_types):
strides = (strides,) * len(kernel_size)
if isinstance(padding, numeric_types):
padding = (padding,) * len(kernel_size)
if isinstance(dilation, numeric_types):
dilation = (dilation,) * len(kernel_size)
self._op_name = op_name

offset_channels = 27
self._kwargs_offset = {
'kernel': kernel_size, 'stride': strides, 'dilate': dilation,
'pad': padding, 'num_filter': offset_channels, 'num_group': groups,
'no_bias': not offset_use_bias, 'layout': layout}

self._kwargs_deformable_conv = {
'kernel': kernel_size, 'stride': strides, 'dilate': dilation,
'pad': padding, 'num_filter': channels, 'num_group': groups,
'num_deformable_group': num_deformable_group,
'no_bias': not use_bias, 'layout': layout}

if adj:
self._kwargs_offset['adj'] = adj
self._kwargs_deformable_conv['adj'] = adj

deformable_conv_weight_shape = [0] * (len(kernel_size) + 2)
deformable_conv_weight_shape[0] = channels
deformable_conv_weight_shape[2] = kernel_size[0]
deformable_conv_weight_shape[3] = kernel_size[1]

self.deformable_conv_weight = self.params.get('deformable_conv_weight',
shape=deformable_conv_weight_shape,
init=weight_initializer,
allow_deferred_init=True)

if use_bias:
self.deformable_conv_bias = self.params.get('deformable_conv_bias', shape=(channels,),
init=bias_initializer,
allow_deferred_init=True)
else:
self.deformable_conv_bias = None

dshape = [0] * (len(kernel_size) + 2)
dshape[layout.find('N')] = 1
dshape[layout.find('C')] = in_channels

op = getattr(symbol, 'Convolution')
offset = op(symbol.var('data', shape=dshape), **self._kwargs_offset)

offsetshapes = offset.infer_shape_partial()[0]

self.offset_weight = self.params.get('offset_weight', shape=offsetshapes[1],
init=offset_weight_initializer,
allow_deferred_init=True)

if offset_use_bias:
self.offset_bias = self.params.get('offset_bias', shape=offsetshapes[2],
init=offset_bias_initializer,
allow_deferred_init=True)
else:
self.offset_bias = None

if activation:
self.act = Activation(activation, prefix=activation + '_')
else:
self.act = None

def hybrid_forward(self, F, x, offset_weight, deformable_conv_weight, offset_bias=None, deformable_conv_bias=None):
if offset_bias is None:
offset = F.Convolution(x, offset_weight, cudnn_off=True, **self._kwargs_offset)
else:
offset = F.Convolution(x, offset_weight, offset_bias, cudnn_off=True, **self._kwargs_offset)

offset_t = F.slice_axis(offset, axis=1, begin=0, end=18)
mask = F.slice_axis(offset, axis=1, begin=18, end=None)
mask = F.sigmoid(mask) * 2

if deformable_conv_bias is None:
act = F.contrib.ModulatedDeformableConvolution(data=x, offset=offset_t, mask=mask,
weight=deformable_conv_weight,
name='fwd', **self._kwargs_deformable_conv)
else:
act = F.contrib.ModulatedDeformableConvolution(data=x, offset=offset_t, mask=mask,
weight=deformable_conv_weight,
bias=deformable_conv_bias, name='fwd',
**self._kwargs_deformable_conv)

if self.act:
act = self.act(act)
return act

def _alias(self):
return 'modulated_deformable_conv'
Loading

0 comments on commit 76e21ff

Please sign in to comment.