pytorch使用sobel、laplace算子

在torch.nn.functional.conv2d()中使用自己构造的sobel卷积核，卷积核shape为（out_channels, in_channels, kernel_size, kernel_size）

import torch, cv2
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

sobel_x = torch.tensor([[-1, -2, -1],
                        [0, 0, 0],
                        [1, 2, 1]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3)
sobel_y = torch.tensor([[-1, 0, 1],
                        [-2, 0, 2],
                        [-1, 0, 1]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3)
laplace = torch.tensor([[0, 1, 0],
                        [1, -4, 1],
                        [0, 1, 0]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3)
avgpool = torch.tensor([[1/9, 1/9, 1/9],
                         [1/9, 1/9, 1/9],
                         [1/9, 1/9, 1/9]], dtype=torch.float, requires_grad=False).view(1, 1, 3, 3)

def conv_operator(filename, kernel, in_channels=1):
    if in_channels == 1:
        img = np.expand_dims(cv2.imread(filename, 0), 2)    # gray
    elif in_channels == 3:
        img = cv2.imread(filename, 1)                        # bgr
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        exit()

    x = torch.from_numpy(img.transpose([2, 0, 1])).unsqueeze(0).float()
    y = F.conv2d(x, kernel.repeat(1, in_channels, 1, 1), stride=1, padding=1,)
    y = y.squeeze(0).numpy().transpose(1, 2, 0)

    return img, y

def plt_show(windowsname, img, channels=1):
    plt.figure(windowsname)
    if channels ==1:
        plt.imshow(img, cmap='gray')
    elif channels == 3:
        plt.imshow(img, )
    else:
        exit()

    plt.axis('on')
    plt.show()

if __name__=="__main__":
    img_name = 'bg.bmp'

    img, y = conv_operator(img_name, sobel_x, 3)
    plt_show("input img", img, 3)
    plt_show("sobel_x", y)

    _, y = conv_operator(img_name, sobel_y, 1)
    plt_show("sobel_y", y)

    _, y = conv_operator(img_name, laplace, 3)
    plt_show("laplace", y)

    _, y = conv_operator(img_name, avgpool, 3)
    plt_show("avgpool", y)

show：