pytorch使用卷积神经网络提取图片边缘信息

import numpy as np
import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt

im=Image.open('./girl.jpg').convert('L')
im = np.array(im, dtype='float32')
plt.figure()
plt.subplot(1,2,1)
plt.imshow(im.astype('uint8'),cmap='gray')

im1 = torch.from_numpy(im.reshape((1,1,im.shape[0],im.shape[1])))
conv1 = nn.Conv2d(1,	1,	3,	bias=False)
sobel_kernel = np.array([[-1,	-1,	-1],	[-1,	8,	-1],	[-1,	-1,	-1]],	dtype='float32')
sobel_kernel = sobel_kernel.reshape((1,	1,	3,	3))
conv1.weight.data = torch.from_numpy(sobel_kernel)
edge1 = conv1(Variable(im1))
edge1 = edge1.data.squeeze().numpy()
plt.subplot(1,2,2)
plt.imshow(edge1,	cmap='gray')
plt.show()