pytorch学习之路——08

卷积神经网络：自定义卷积核、查看卷积操作后的特征图

import cv2
import numpy as np
import matplotlib.pyplot as plt

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


# 读数据并转换成灰度图
img_path = "./data/tenyounger.jpg"
bgr_img = cv2.imread(img_path)
gray_img = cv2.cvtColor(bgr_img,cv2.COLOR_BGR2GRAY)
gray_img = gray_img.astype("float32")/255
plt.imshow(gray_img,cmap="gray")
plt.show()

# 定义卷积
filter_vals = np.array([[-1,-1,1,1],[-1,-1,1,1],[-1,-1,1,1],[-1,-1,1,1]])
print("Filter shape: ",filter_vals.shape)
filter1 = filter_vals
filter2 = -filter1
filter3 = filter1.T
filter4 = -filter3
filters = np.array([filter1, filter2, filter3, filter4])
print("Filter1: \n",filter1)

# 可视化卷积核
def viz_conv(filters):
    fig = plt.figure(figsize=(5,5))
    for i in range(4):
        ax = fig.add_subplot(1, 4, i+1, xticks=[], yticks=[])
        ax.imshow(filters[i], cmap='gray')
        ax.set_title('Filter %s' % str(i+1))
        width, height = filters[i].shape
        for x in range(width):
            for y in range(height):
                ax.annotate(str(filters[i][x][y]), xy=(y,x),
                            horizontalalignment='center',
                            verticalalignment='center',
                            color='white' if filters[i][x][y]<0 else 'black')
    plt.show()

# 可视化卷积后的输出
def viz_layer(layer, n_filters=4):
    fig = plt.figure(figsize=(8,8))
    for i in range(n_filters):
        ax = fig.add_subplot(1, n_filters, i + 1, xticks=[], yticks=[])
        # grab layer outputs
        ax.imshow(np.squeeze(layer[0, i].data.numpy()), cmap='gray')
        ax.set_title('Output %s' % str(i + 1))

# 定义卷积层
class Net(nn.Module):
    def __init__(self,weight):
        super().__init__()
        k_height,k_width = weight.shape[2:]
        self.conv = nn.Conv2d(1,4,kernel_size=(k_height,k_width),bias=False)
        # 使用自定义的卷积
        self.conv.weight = torch.nn.Parameter(weight)
        self.pool = nn.MaxPool2d(2,2)
        self.norm = nn.BatchNorm2d(4)

    def forward(self,x):
        conv_x = self.conv(x)
        act_x = F.relu(conv_x)
        pooled_x = self.pool(act_x)
        norm_x = self.norm(pooled_x)

        return conv_x,act_x,pooled_x,norm_x

weight = torch.from_numpy(filters).unsqueeze(1).type(torch.FloatTensor)
model = Net(weight)

gray_img_tensor = torch.from_numpy(gray_img).unsqueeze(0).unsqueeze(1)
conv_layer,act_layer,pool_layer,norm_layer = model.forward(gray_img_tensor)
print(conv_layer.shape)
viz_layer(conv_layer)
viz_layer(act_layer)
viz_layer(pool_layer)
viz_layer(norm_layer)
plt.show()