Pytorch 可视化CNN中间层的输出

可视化中间层的输出有助于理解输入图像如何在不同层之间进行转换。

PyTorch 提供了一个名为register_forward_hook的方法，它允许传入一个可以提取特定层输出的函数。

以下例子，使用VGG-16网络和猫狗图像分类，演示如何提取不同层的输出。

import matplotlib.pyplot as plt
import torch
from torchvision import models, transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import os
import numpy as np
from torchvision.datasets import ImageFolder

torch.cuda.set_device(0)  # 设置GPU ID
is_cuda = True
simple_transform = transforms.Compose([transforms.Resize((224, 224)),
                                       transforms.ToTensor(),  # H, W, C -> C, W, H 归一化到(0,1)，简单直接除以255
                                       transforms.Normalize([0.485, 0.456, 0.406],  # std
                                                            [0.229, 0.224, 0.225])])

# mean  先将输入归一化到(0,1)，再使用公式”(x-mean)/std”，将每个元素分布到(-1,1)
# 使用 ImageFolder 必须有对应的目录结构
train = ImageFolder("./datas/dogs-vs-cats/train", simple_transform)
valid = ImageFolder("./datas/dogs-vs-cats/valid", simple_transform)
train_loader = DataLoader(train, batch_size=1,  shuffle=False, num_workers=5)
val_loader = DataLoader(valid, batch_size=1, shuffle=False, num_workers=5)

vgg = models.vgg16(pretrained=True).cuda()

# 提取不同层输出的 主要代码
class LayerActivations:
    features = None

    def __init__(self, model, layer_num):
        self.hook = model[layer_num].register_forward_hook(self.hook_fn)

    def hook_fn(self, module, input, output):
        self.features = output.cpu()

    def remove(self):
        self.hook.remove()

print(vgg.features)

conv_out = LayerActivations(vgg.features, 0)  # 提出第 一个卷积层的输出 
img = next(iter(train_loader))[0]

# imshow(img)
o = vgg(Variable(img.cuda()))
conv_out.remove()  #
act = conv_out.features  # act 即 第0层输出的特征

# 可视化 输出
fig = plt.figure(figsize=(20, 50))
fig.subplots_adjust(left=0, right=1, bottom=0, top=0.8, hspace=0, wspace=0.2)
for i in range(30):
    ax = fig.add_subplot(12, 5, i+1, xticks=[], yticks=[])
    ax.imshow(act[0][i].detach().numpy(), cmap="gray")

plt.show()

原图像：

第一个VGG-16 卷积层输出 部分结果可视化，如下图：

可视化 VGG-16最后的输出结果，设置：（29 是最后一个MaxPool2d的ID）

conv_out = LayerActivations(vgg.features, 29)

可视化结果， 如下图：

 

 从上可以 看出， 低层的输出的特征 人还可以看明白， 而更高层的输出 倾向于学习更高层次的特征，而解释性比较差。

可视化层的输出， 对于理解CNN有很大的帮助，神经网络如果一直是一个黑盒，人类使用起来恐怕也不放心。解释神经网络是如何学习的，必然是一个研究方向。