PyTorch nn.Conv2d自定义权重实现三通道上分通道滤波

PyTorch nn.Conv2d自定义权重实现三通道上分通道滤波

最近做的一个项目涉及到传统图像滤波算法与深度学习算法结合的框架，为了实现模型的统一，考虑用PyTorch重写之前用OpenCV实现的一种MSRCR算法的变种，该算法中需要用到一种自定义权重的滤波算子来代替高斯滤波。本文不涉及算法的具体实现，重点在于如何实现自定义权重的nn.Conv2d层。

输入

RGB（如用OpenCV读取则为BGR，通道先后顺序不影响结果）三通道图像，从ndarray转为dtype=torch.float32，为了避免算法中torch.log10(data)中运算数为0带来的数值不稳定，输入范围变为[1.0, 256.0]。
注意：nn.Conv2d的输入维度必须为四维，(batch_size, channel, height, width)，这里有三个用于维度调整的方法

torch.tensor().unsqueeze(0)      # 在维度0上添加一位，即shape变换：(3, 3) -> (1, 3, 3)
torch.tensor().squeeze(0)        # 在维度0上缩减一位，即shape变换：(1, 3, 3) -> (3, 3)
torch.tensor().repeat(a, b, c, d) # 复制维度，如(1, 1, 3, 3).repeat(3, 3, 3, 3) -> (3, 3, 9, 9)

输入在卷积前先做ReflectionPadding：

pad = nn.ReflectionPad2d(int((k_size - 1) / 2))
x = pad(x)

边缘镜像填充，这种方法与OpenCV中cv2.filter2D的默认填充方式cv2.BORDER_REFLECT_101相似。

nn.Conv2d构建

根据Conv2d的PyTorch官方文档，nn.Conv2d构建参数如下：

torch.nn.Conv2d(
	in_channels,
	out_channels, 
	kernel_size, 
	stride=1, 
	padding=0, 
	dilation=1, 
	groups=1,
	bias=True, 
	padding_mode='zeros', 
	device=None, 
	dtype=None)

网上常见的做法如下：
假设卷积核k_size = 3，先定义一个nn.Conv2d层，将模型移动到GPU上（如果没有则忽略最后的.to(self.device)）

conv = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=k_size, padding=0, bias=False).to(self.device)

获取一个自定义的卷积核

filter = self.my_filter(kernel_size=k_size)

直接修改卷积层的参数

conv.weight.data = torch.tensor(filter, dtype=torch.float32).unsqueeze(0).unsqueeze(0).repeat(1, 1, 1, 1).to(self.device)

滤波

out = conv(x)

此时conv.weight.data.shape == (1， 1， 3， 3)。但是，上述参数只能应用于单通道图。

三通道上应用nn.Conv2d

为了实现基础的三通道上的滤波，这里首先将第一步改成

conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=k_size, padding=0, bias=False).to(self.device)

接下来

conv.weight.data = torch.tensor(filter, dtype=torch.float32).unsqueeze(0).unsqueeze(0).repeat(3, 3, 1, 1).to(self.device)
out = conv(x)

这样是错误的！
对于矩阵

[[10, 12, 10],
[11, 13, 11],
[10, 12, 10]]

在卷积核

[[0.625, 1.5, 0.625],
[1.375, 3.25, 1.375],
[0.625, 1.5, 0.625]]

下，卷积得到的值是29.5，而OpenCV将得到11.5（均不考虑边缘填充）。
注意到得到的结果约为正确结果的三倍，考虑卷积神经网络的实现：

左侧黄色为输入三通道矩阵，右侧为输出三通道矩阵，如果未指定卷积的方式，得到的结果应该是类似于全连接网络，对每一个输入通道都用每一个卷积核进行一次卷积，再加权求和。而我们希望的形式应该如下：

为此，注意到nn.Conv2d有一个group参数用于分组卷积：

机翻版本如下：

也就是说，当groups = 3时，每个输入通道都与自己的一组卷积核做卷积，这一组卷积核的大小是out_channels/in_channels=1，也就是说实现了前文图2的功能。
groups参数在深度可分离卷积的实现中也有使用。

实现

前文代码改为

conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=k_size, padding=0, groups=3, bias=False).to(self.device)

获取一个自定义的卷积核

filter = self.my_filter(kernel_size=k_size)

修改卷积层的参数，此处repeat的次数我没有深究，总之最后两个维度的结果应为(3, 3)，而前两个维度这样写不报错即可，猜测第一个参数是通道数，第二个参数是每组的卷积核个数

conv.weight.data = torch.tensor(filter, dtype=torch.float32).unsqueeze(0).unsqueeze(0).repeat(3, 1, 1, 1).to(self.device)

最后卷积

out = conv(x)