Spaces:
Runtime error
Runtime error
File size: 2,912 Bytes
68c537d |
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 |
import math
import torch
import torch.nn as nn
from torch.nn.modules.utils import _pair
from mmcv.ops import modulated_deform_conv2d
class DCN_layer_rgb(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1,
groups=1, deformable_groups=1, bias=True, extra_offset_mask=True):
super(DCN_layer_rgb, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = _pair(kernel_size)
self.stride = stride
self.padding = padding
self.dilation = dilation
self.groups = groups
self.deformable_groups = deformable_groups
self.with_bias = bias
self.weight = nn.Parameter(
torch.Tensor(out_channels, in_channels, *self.kernel_size))
self.extra_offset_mask = extra_offset_mask
self.conv_offset_mask = nn.Conv2d(
self.in_channels,
self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1],
kernel_size=self.kernel_size, stride=_pair(self.stride), padding=_pair(self.padding),
bias=True
)
self.c1 = nn.Conv2d(in_channels*4, out_channels, 1, 1, 0, bias=False)
self.c2 = nn.Conv2d(out_channels, out_channels, 1, 1, 0, bias=False)
if bias:
self.bias = nn.Parameter(torch.Tensor(out_channels))
else:
self.register_parameter('bias', None)
self.init_offset()
self.reset_parameters()
def reset_parameters(self):
n = self.in_channels
for k in self.kernel_size:
n *= k
stdv = 1. / math.sqrt(n)
self.weight.data.uniform_(-stdv, stdv)
if self.bias is not None:
self.bias.data.zero_()
def init_offset(self):
self.conv_offset_mask.weight.data.zero_()
self.conv_offset_mask.bias.data.zero_()
def forward(self, input_feat, inter, fea):
b, c, h, w = input_feat.shape
fea = self.c1(fea).unsqueeze(1)
weight = self.weight.unsqueeze(0) * fea
weight = weight.view(b * self.out_channels, self.in_channels, self.kernel_size[0],
self.kernel_size[1]).contiguous()
input_feat = input_feat.view(1, b * self.in_channels, h, w)
out = self.conv_offset_mask(inter)
o1, o2, mask = torch.chunk(out, 3, dim=1)
offset = torch.cat((o1, o2), dim=1)
mask = torch.sigmoid(mask)
out = modulated_deform_conv2d(input_feat.contiguous(), offset, mask, weight, self.bias, self.stride,
self.padding, self.dilation, b, b)
_, _, height, width = out.shape
out = out.view(b, self.out_channels, height, width).contiguous()
out2 = self.c2(out)
return out2
|