如何在 pytorch 中为 Conv2d 实现棋盘格步幅？

Question

我正在尝试使用 pytorch 创建一个 convnet 来处理二维矩阵的输入。 我正在使用 3x5 过滤器，我希望它具有如下自定义步幅 - 在偶数行号上，我希望过滤器从位置 0（图像中的红色）处的元素开始，在奇数行号上我希望它开始在位置 1 的元素上（图像中的蓝色），并且在这两种情况下，x 方向上的步幅均为 2。 这意味着如果我将图像中的矩阵作为我的输入，我希望过滤器的中心只有 0。 我知道这在 convnets 中非常不寻常，但这实际上是物理学中的一个问题，因此确切的步幅很重要。

Answer 1

以下自定义 conv2d 层以棋盘格步幅实现卷积，如原始问题所示。 这里的困难在于pytorch并不真正支持这样的不一致步幅。 也就是说，我们可以将此操作分解为两个单独的跨步卷积，一个用于偶数行，另一个用于奇数行。 之后，我们可以将结果重新交织在一起。 下面的代码中有一些细节可以确保我们正确填充（如果需要）。 此外，该层完全支持反向传播。

import torch.nn as nn
import torch.nn.functional as F

class AMNI_Conv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding=0, bias=True):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=bias, padding=padding)
        self.crow = self.conv.kernel_size[0] // 2
        self.ccol = self.conv.kernel_size[1] // 2

        # this module only works with odd sized kernels
        assert self.conv.kernel_size[0] % 2 == 1 and self.conv.kernel_size[1] % 2 == 1

    def forward(self, x):
        # currently only padding with zeros is supported
        if self.conv.padding[0] != 0 or self.conv.padding[1] != 0:
            x = F.pad(x, pad=(self.conv.padding[1], self.conv.padding[1], self.conv.padding[0], self.conv.padding[0]))

        # center filters on the "zeros" according to the diagram by AMNI, starting column for even/odd rows may need to change depending on padding/kernel size
        if (self.crow + self.ccol + self.conv.padding[0] + self.conv.padding[1]) % 2 == 0:
            x_even = F.conv2d(x[:, :, :-1, 1:], self.conv.weight, self.conv.bias, stride=2)
            x_odd = F.conv2d(x[:, :, 1:, :-1], self.conv.weight, self.conv.bias, stride=2)
        else:
            x_even = F.conv2d(x[:, :, :-1, :-1], self.conv.weight, self.conv.bias, stride=2)
            x_odd = F.conv2d(x[:, :, 1:, 1:], self.conv.weight, self.conv.bias, stride=2)
        b, c, h, w = x_even.shape

        # interleave even and odd rows back together
        return torch.stack((x_even, x_odd), dim=3).contiguous().view(b, c, -1, w)

---

例子

这一层基本上就像一个普通的 Conv2d，但具有棋盘格步幅。

>>> x = torch.arange(64).view(1, 1, 8, 8).float()
tensor([[[[ 0.,  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.]]]])

---

>>> layer = AMNI_Conv2d(1, 1, (3, 5), bias=False)

# set kernels to delta functions to demonstrate kernel centers
>>> with torch.no_grad():
...     layer.conv.weight.zero_()
...     layer.conv.weight[:,:,1,2] = 1

>>> result = layer(x)
tensor([[[[10., 12.],
          [19., 21.],
          [26., 28.],
          [35., 37.],
          [42., 44.],
          [51., 53.]]]], grad_fn=<ViewBackward>)

您也可以使用填充来获取原始图中的每个“零”

>>> layer = AMNI_Conv2d(1, 1, (3, 5), padding=(1, 2), bias=False)

# set kernels to delta functions to demonstrate kernel centers
>>> with torch.no_grad():
...     layer.conv.weight.zero_()
...     layer.conv.weight[:,:,1,2] = 1

>>> result = layer(x)
tensor([[[[ 1.,  3.,  5.,  7.],
          [ 8., 10., 12., 14.],
          [17., 19., 21., 23.],
          [24., 26., 28., 30.],
          [33., 35., 37., 39.],
          [40., 42., 44., 46.],
          [49., 51., 53., 55.],
          [56., 58., 60., 62.]]]], grad_fn=<ViewBackward>)