-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
deformable_convolution_2d.py
133 lines (109 loc) · 5.32 KB
/
deformable_convolution_2d.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
from chainer.functions import deformable_convolution_2d_sampler
from chainer import initializers
from chainer.initializers import constant
from chainer import link
from chainer.links.connection.convolution_2d import Convolution2D
from chainer import variable
class DeformableConvolution2D(link.Chain):
"""Two-dimensional deformable convolutional layer.
This link wraps the
convolution layer for offset prediction and
the :func:`~chainer.functions.deformable_convolution_2d_sampler`
function.
This also holds the filter weights and bias vectors of two
convolution layers as parameters.
Args:
in_channels (int): Number of channels of input arrays. If ``None``,
parameter initialization will be deferred until the first forward
data pass at which time the size will be determined.
out_channels (int): Number of channels of output arrays.
ksize (int or pair of ints): Size of filters (a.k.a. kernels).
``ksize=k`` and ``ksize=(k, k)`` are equivalent.
stride (int or pair of ints): Stride of filter applications.
``stride=s`` and ``stride=(s, s)`` are equivalent.
pad (int or pair of ints): Spatial padding width for input arrays.
``pad=p`` and ``pad=(p, p)`` are equivalent.
offset_nobias (bool): If ``True``, then this link does not use the
bias term for the first convolution layer.
offset_initialW (:ref:`initializer <initializer>`): Initializer to
initialize the weight of the first convolution layer.
When it is :class:`numpy.ndarray`, its ``ndim`` should be 4.
offset_initial_bias (:ref:`initializer <initializer>`): Initializer to
initialize the bias of the first convolution layer.
If ``None``, the bias will be initialized to
zero. When it is :class:`numpy.ndarray`, its ``ndim`` should be 1.
deform_nobias (bool): If ``True``, then this link does not use the
bias term for the second convolution layer.
deform_initialW (:ref:`initializer <initializer>`): Initializer to
initialize the weight for the second convolution layer.
When it is :class:`numpy.ndarray`,
its ``ndim`` should be 4.
deform_initial_bias (:ref:`initializer <initializer>`): Initializer to
initialize the bias for the second convolution layer.
If ``None``, the bias will be initialized to
zero. When it is :class:`numpy.ndarray`, its ``ndim`` should be 1.
.. seealso::
See :func:`chainer.functions.deformable_convolution_2d_sampler`.
"""
def __init__(self, in_channels, out_channels, ksize, stride=1, pad=0,
offset_nobias=False, offset_initialW=None,
offset_initial_bias=None,
deform_nobias=False,
deform_initialW=None, deform_initial_bias=None):
super(DeformableConvolution2D, self).__init__()
kh, kw = _pair(ksize)
with self.init_scope():
self.offset_conv = Convolution2D(
in_channels, 2 * kh * kw, ksize, stride, pad,
offset_nobias, offset_initialW, offset_initial_bias)
self.deform_conv = DeformableConvolution2DSampler(
in_channels, out_channels, ksize, stride, pad,
deform_nobias, deform_initialW, deform_initial_bias)
def forward(self, x):
"""Applies the deformable convolution.
Args:
x (~chainer.Variable): Input image.
Returns:
~chainer.Variable: Output of the deformable convolution.
"""
offset = self.offset_conv(x)
return self.deform_conv(x, offset)
class DeformableConvolution2DSampler(link.Link):
"""Apply a two-dimensional deformable convolution layer using offsets"""
def __init__(self, in_channels, out_channels, ksize, stride=1, pad=0,
nobias=False, initialW=None, initial_bias=None):
super(DeformableConvolution2DSampler, self).__init__()
self.ksize = ksize
self.stride = _pair(stride)
self.pad = _pair(pad)
self.out_channels = out_channels
self.initialW = initialW
if initialW is None:
initialW = constant.Zero()
with self.init_scope():
W_initializer = initializers._get_initializer(initialW)
self.W = variable.Parameter(W_initializer)
if nobias:
self.b = None
else:
if initial_bias is None:
initial_bias = initializers.Constant(0)
bias_initializer = initializers._get_initializer(initial_bias)
self.b = variable.Parameter(bias_initializer)
if in_channels is not None:
self._initialize_params(in_channels)
def _initialize_params(self, in_channels):
kh, kw = _pair(self.ksize)
W_shape = (self.out_channels, in_channels, kh, kw)
self.W.initialize(W_shape)
if self.b is not None:
self.b.initialize(self.out_channels)
def forward(self, x, offset):
if self.W.array is None:
self._initialize_params(x.shape[1])
return deformable_convolution_2d_sampler(
x, offset, self.W, self.b, self.stride, self.pad)
def _pair(x):
if hasattr(x, '__getitem__'):
return x
return x, x