|
|
|
import math |
|
|
|
from torch import nn |
|
from torch.nn import functional as F |
|
|
|
from .registry import CONV_LAYERS |
|
|
|
|
|
@CONV_LAYERS.register_module() |
|
class Conv2dAdaptivePadding(nn.Conv2d): |
|
"""Implementation of 2D convolution in tensorflow with `padding` as "same", |
|
which applies padding to input (if needed) so that input image gets fully |
|
covered by filter and stride you specified. For stride 1, this will ensure |
|
that output image size is same as input. For stride of 2, output dimensions |
|
will be half, for example. |
|
|
|
Args: |
|
in_channels (int): Number of channels in the input image |
|
out_channels (int): Number of channels produced by the convolution |
|
kernel_size (int or tuple): Size of the convolving kernel |
|
stride (int or tuple, optional): Stride of the convolution. Default: 1 |
|
padding (int or tuple, optional): Zero-padding added to both sides of |
|
the input. Default: 0 |
|
dilation (int or tuple, optional): Spacing between kernel elements. |
|
Default: 1 |
|
groups (int, optional): Number of blocked connections from input |
|
channels to output channels. Default: 1 |
|
bias (bool, optional): If ``True``, adds a learnable bias to the |
|
output. Default: ``True`` |
|
""" |
|
|
|
def __init__(self, |
|
in_channels, |
|
out_channels, |
|
kernel_size, |
|
stride=1, |
|
padding=0, |
|
dilation=1, |
|
groups=1, |
|
bias=True): |
|
super().__init__(in_channels, out_channels, kernel_size, stride, 0, |
|
dilation, groups, bias) |
|
|
|
def forward(self, x): |
|
img_h, img_w = x.size()[-2:] |
|
kernel_h, kernel_w = self.weight.size()[-2:] |
|
stride_h, stride_w = self.stride |
|
output_h = math.ceil(img_h / stride_h) |
|
output_w = math.ceil(img_w / stride_w) |
|
pad_h = ( |
|
max((output_h - 1) * self.stride[0] + |
|
(kernel_h - 1) * self.dilation[0] + 1 - img_h, 0)) |
|
pad_w = ( |
|
max((output_w - 1) * self.stride[1] + |
|
(kernel_w - 1) * self.dilation[1] + 1 - img_w, 0)) |
|
if pad_h > 0 or pad_w > 0: |
|
x = F.pad(x, [ |
|
pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2 |
|
]) |
|
return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, |
|
self.dilation, self.groups) |
|
|