67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torchvision.ops import deform_conv2d
|
|
|
|
|
|
class DeformableConv2d(nn.Module):
|
|
def __init__(self,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
padding=1,
|
|
bias=False):
|
|
|
|
super(DeformableConv2d, self).__init__()
|
|
|
|
assert type(kernel_size) == tuple or type(kernel_size) == int
|
|
|
|
kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size)
|
|
self.stride = stride if type(stride) == tuple else (stride, stride)
|
|
self.padding = padding
|
|
|
|
self.offset_conv = nn.Conv2d(in_channels,
|
|
2 * kernel_size[0] * kernel_size[1],
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=self.padding,
|
|
bias=True)
|
|
|
|
nn.init.constant_(self.offset_conv.weight, 0.)
|
|
nn.init.constant_(self.offset_conv.bias, 0.)
|
|
|
|
self.modulator_conv = nn.Conv2d(in_channels,
|
|
1 * kernel_size[0] * kernel_size[1],
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=self.padding,
|
|
bias=True)
|
|
|
|
nn.init.constant_(self.modulator_conv.weight, 0.)
|
|
nn.init.constant_(self.modulator_conv.bias, 0.)
|
|
|
|
self.regular_conv = nn.Conv2d(in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=self.padding,
|
|
bias=bias)
|
|
|
|
def forward(self, x):
|
|
#h, w = x.shape[2:]
|
|
#max_offset = max(h, w)/4.
|
|
|
|
offset = self.offset_conv(x)#.clamp(-max_offset, max_offset)
|
|
modulator = 2. * torch.sigmoid(self.modulator_conv(x))
|
|
|
|
x = deform_conv2d(
|
|
input=x,
|
|
offset=offset,
|
|
weight=self.regular_conv.weight,
|
|
bias=self.regular_conv.bias,
|
|
padding=self.padding,
|
|
mask=modulator,
|
|
stride=self.stride,
|
|
)
|
|
return x
|