From 4d958ebbaa6fc7aab3a6d06396181d929e1db7fa Mon Sep 17 00:00:00 2001 From: Yonghye Kwon Date: Tue, 16 Nov 2021 08:08:39 +0900 Subject: [PATCH] support arbitrary kernel_size fix #10 --- dcn.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/dcn.py b/dcn.py index 2b651f2..4fe0f67 100644 --- a/dcn.py +++ b/dcn.py @@ -13,11 +13,14 @@ def __init__(self, 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 * kernel_size, + 2 * kernel_size[0] * kernel_size[1], kernel_size=kernel_size, stride=stride, padding=self.padding, @@ -27,7 +30,7 @@ def __init__(self, nn.init.constant_(self.offset_conv.bias, 0.) self.modulator_conv = nn.Conv2d(in_channels, - 1 * kernel_size * kernel_size, + 1 * kernel_size[0] * kernel_size[1], kernel_size=kernel_size, stride=stride, padding=self.padding, @@ -59,3 +62,4 @@ def forward(self, x): stride=self.stride, ) return x +