一分钟教你在PyTorch跑模型的时候提取中间层查看图片

首先导入需要的库

import torch
import torch.nn as nn  # 网络库
import torch.nn.functional as F
import torch.optim as optim  # 优化器

import pandas as pd  # 数据处理

import torchvision
import torchvision.transforms as transforms   

torch.set_grad_enabled(True)
torch.set_printoptions(linewidth=120)

from torch.utils.data import DataLoader  # 数据库
from IPython.display import display, clear_output  # 可以在notebook版本中显示pandas

from torchvision import transforms
from PIL import Image
import cv2

然后搭建模型，我们这里将cnn从第一层分开了，只留下了第一层的卷积层，也就是我们要查看第一层卷积后的样子

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)  # 定义卷积层
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)

        self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)  # 定义线性层
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)  # 定义输出层
    
    # 开始搭建
    def forward(self, t):
        # (1) 输入层，其实可以不写，但写成这样更清楚
        t = t

        # (2) 第一个卷积层，卷积层后接激励函数再接最大池化层
        t = self.conv1(t)
        
        # t = F.relu(t)
        # t = F.max_pool2d(t, kernel_size=2, stride=2)

        # # (3) 第二个卷积层，卷积层后接激励函数再接最大池化层
        # t = self.conv2(t)
        # t = F.relu(t)
        # t = F.max_pool2d(t, kernel_size=2, stride=2)

        # # (4) 第一个线性层，进线性层之前先reshape成1*192(12*4*4)的维度，reshape函数中的-1代表n，根据
        # t = t.reshape(-1, 12 * 4 *4)
        # print('after reshaped:')
        # print(t.shape)
        # t = self.fc1(t)
        # t = F.relu(t)

        # # (5) 第二个线性层
        # t = self.fc2(t)
        # t = F.relu(t)

        # # (6) 输出层
        # t = self.out(t)

        return t

# 获取预测正确的个数
def get_num_correct(preds, labels):
    return preds.argmax(dim=1).eq(labels).sum().item()

接下来加载数据

train_set = torchvision.datasets.FashionMNIST(
    root='./data/FashionMNIST'
    ,train=True
    ,download=True
    ,transform=transforms.Compose([
        transforms.ToTensor()
    ])
)

train_loader = torch.utils.data.DataLoader(train_set, batch_size=1)
batch = next(iter(train_loader))
images, labels = batch

创建模型，并将图片输入进去返回一个tensor，这个tensor就是我们得到的卷积后的六个图片（因为我们输出通道为6）

net = Network()
t = net(images)

把第一维度的bach_size给去掉

print(t.shape)
t = t.squeeze(0)
print(t.shape)  # 此处应该是[6, 24, 24]

这一步就是显示出所有图片，自己根据你的通道数改plt.subplot函数的前两个参数，我这边是6张图片，所有我就绘制了2*3的图片矩阵，效果如下图：

from matplotlib.pyplot import imshow
import matplotlib.pyplot as plt 
import numpy as np
from PIL import Image
%matplotlib inline
#Tensor转成PIL.Image重新显示
for i in range(1, len(t) + 1):
    plt.subplot(2,3,i), plt.title('imge:oringe')
    new_img_PIL = transforms.ToPILImage()(t[i-1]).convert()
    imshow(np.asarray(new_img_PIL))

我是在notebook中执行的，如果在pycharm中可能需要删掉inline这行