pytorch 自定义卷积核进行卷积--卷积核订制

镜像问题：tensorflow 卷积：设定特定卷积核

动机：

有时候我们希望对神经网络得到的张量进行特定的卷积操作，如求梯度，求边缘等。

或者，有时我们希望把图像进行预处理操作后输入神经网络，而这个预处理可以用卷积操作实现，如求梯度或高斯滤波。我们希望这个卷积操作在GPU上执行，这样可以加快运行速度。

这就需要我们在pytorch中进行自定义的卷积核卷积。

困难：

在常用的pytorch的卷积操作中，只有卷积核的尺寸和通道信息，且定义的卷积核是可训练的变量，无法实现卷积核订制，且保持不变

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True)

解决方法：

幸运的是，torch.nn.functional.conv2d 可以实现卷积核订制。

torch.nn.functional.conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1)

该函数的参数如下：

• input – input tensor of shape (minibatch,in_channels,iH,iW) # [N,C,H,W]

• weight – filters of shape(out_channels,in_channels​,kH,kW) # [C_out,C_in,kH,kW]

• bias – optional bias tensor of shape (out_channels) . Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sH, sW). Default: 1

• padding – implicit paddings on both sides of the input. Can be a single number or a tuple (padH, padW). Default: 0

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dH, dW). Default: 1

• groups – split input into groups, \text{in\_channels}in_channels should be divisible by the number of groups. Default: 1

 

注意事项：

在torch.nn.functional.conv2d中，input.shape=[B,C,H,W],卷积核weight.shape=[out_channels,in_channels,h,w]。这与tensorflow中的习惯不同。

如果输入的数据的shape=[B,C,H,W]，可使用 tensor.permute()操作进行张量维度的重新排列。

应用示例：

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

import numpy as np
import matplotlib.pyplot as plt

def gradient_blur_bhwc(input_bhwc):
	"""
	:param     input_bhwc: shape=[B,H,W,C]
	:return:   grad_x_bhwc: x方向梯度图，shape = [B,H,W-1,1]
		       grad_y_bhwc：y方向梯度图，shape = [B,H-1,W,1]
	"""
	in_channel = input_bhwc.size()[-1]
	input_bchw = input_bhwc.permute(0,3,1,2)

	filter_x = torch.tensor([[[[-1., 1.]]] * in_channel], dtype=torch.float32)  # [outchannel,in_channel,h,w]=>[1,3,1,2]
	filter_y = torch.tensor([[[[-1.], [1.]]] * in_channel],dtype=torch.float32)  # [outchannel,in_channel,h,w]=>[1,3,2,1]
	
	grad_x_bhwc = F.conv2d(input_bchw, weight=filter_x, stride=1).permute(0,2,3,1)
	grad_y_bhwc = F.conv2d(input_bchw, weight=filter_y, stride=1).permute(0,2,3,1)
	
	return grad_x_bhwc,grad_y_bhwc


def gradient_blur_bchw(input_bchw):
	"""
	:param      input_bchw: shape=[B,C,H,W]
	:return:	grad_x_bchw: x方向梯度图，shape = [B,1,H,W-1]
		        grad_y_bchw：y方向梯度图，shape = [B,1,H-1,W]
	"""
	in_channel = input_bchw.size()[1]
	
	filter_x = torch.tensor([[[[-1., 1.]]] * in_channel], dtype=torch.float32)  # [outchannel,in_channel,h,w]=>[1,3,1,2]
	filter_y = torch.tensor([[[[-1.], [1.]]] * in_channel],dtype=torch.float32)  # [outchannel,in_channel,h,w]=>[1,3,2,1]
	
	grad_x_bchw = F.conv2d(input_bchw, weight=filter_x, stride=1)
	grad_y_bchw = F.conv2d(input_bchw, weight=filter_y, stride=1)
	
	return grad_x_bchw, grad_y_bchw

input_bhwc = torch.tensor(np.expand_dims(cv2.imread("001.jpg"),axis=0),dtype=torch.float32) # b,h,w,c
input_bchw = input_bhwc.permute(0,3,1,2)
grad_x_bchw, grad_y_bchw = gradient_blur_bchw(input_bchw)

# 可视化
print(grad_x_bchw.shape,grad_y_bchw.shape)
plt.figure()
plt.subplot(131)
plt.imshow(input_bhwc[0,...]/255.,cmap='gray')
plt.subplot(132)
plt.imshow(grad_x_bchw[0,0,...],cmap='gray')
plt.subplot(133)
plt.imshow(grad_y_bchw[0,0,...],cmap='gray')
plt.show()

参考：pytorch 自定义高斯核进行卷积操作