Pytorch可视化特征图

为什么要可视化特征图？

虽然我们经常讲神经网络是一个黑盒子，不具备解释性，但是人们仍然努力地想要打开这个黑盒子。通过可视化特征图可以让我们较为直观地认识到神经网络到底学到了什么东西（配以卷积核参数更佳），尽管这种认识基本都是经验性的，但也常常可以给人以灵感，促使人们改进算法。

比如其中的代表作之一，2013年Zeiler和Fergus发表的《Visualizing and Understanding Convolutional Networks》

早期LeCun 1998年的文章《Gradient-Based Learning Applied to Document Recognition》中的一张图也非常精彩，个人觉得比Zeiler 2013年的文章更能给人以启发。从下图的F6特征，我们可以清楚地看到原始书写差异非常大的图片如何在深层特征中体现出其不变性。

接下来我们看看如何在PyTorch中可视化特征图。

PyTorch中与可视化特征图相关的函数/接口

因为PyTorch使用动态图机制，所以我们可以在运行过程中将特征图（Tensor）拿出来，只要能得到特征图，基本上就成功了99%，剩下的画图工作虽然可能有一些麻烦，但也只是麻烦而已，构不成什么障碍。

拿到特征图的方法有多种，有人可以从输入开始，一个一个算子地让网络做前向运算，直到想要的特征图处将其返回，这种方法尽管也可行，但略有些麻烦。实际上PyTorch给了一个专用接口可以在前向过程中获取到特征图，这个接口是torch.nn.Module.register_forward_hook。当我们拿到特征图后，PyTorch又有专门的画图和保存图片的接口：torchvision.utils.make_grid和torchvision.utils.save_image，非常方便。

torch.nn.Module.register_forward_hook

函数如其名，这是一个钩子函数，经设置后，会在网络前向的过程中顺带执行。

函数声明：register_forward_pre_hook(hook: Callable[..., None])，括号中的参数是一个函数名，姑且称之为hook_func，hook_func需要自己实现。

hook_func可从前向过程中接收到三个参数：hook_func(module, input, output)。其中module指的是模块的名称，比如对于ReLU模块，module是ReLU()，对于卷积模块，module是Conv2d(in_channel=...)，注意module带有具体的参数。input和output就是我们心心念的特征图，这二者分别是module的输入和输出，输入可能有多个（比如concate层就有多个输入），输出只有一个，所以input是一个tuple，其中每一个元素都是一个Tensor，而输出就是一个Tensor。一般而言output可能更经常拿来做分析。我们可以在hook_func中将特征图画出来并保存为图片，所以hook_func就是我们实现可视化的关键。

torchvision.utils.make_grid, torchvision.utils.save_image

在拿到特征图后，我们其实也不用自己写函数来做画图工作，PyTorch很贴心地为大家准备了两个函数，可以通过这两个函数非常容易地将特征图画出来。两个函数都能够将特征图拼接成一个网格状排列的大图，但两者的执行结果不一样。make_grid将拼接结果作为Tensor返回，save_image直接将拼接结果保存称为磁盘上的图像文件。

这两个函数在使用时特别需要注意将batch_size设为1，此时传入hook_func的input和output的shape就是[1, C, H, W]，需要将shape调整为[C, 1, H, W]后再使用这两个函数，否则会报错或者出现很诡异的问题（不知道是我使用的姿势不对，还是PyTorch自己的bug）。

示例：MNIST

# -*- coding: utf-8 -*-
import os
import shutil
import time
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
import torchvision.utils as vutil
from torch.utils.data import DataLoader
import torchsummary

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
EPOCH = 1
LR = 0.001
TRAIN_BATCH_SIZE = 64
TEST_BATCH_SIZE = 32
BASE_CHANNEL = 32
INPUT_CHANNEL = 1
INPUT_SIZE = 28
MODEL_FOLDER = './save_model'
IMAGE_FOLDER = './save_image'
INSTANCE_FOLDER = None


class Model(nn.Module):
    def __init__(self, input_ch, num_classes, base_ch):
        super(Model, self).__init__()

        self.num_classes = num_classes
        self.base_ch = base_ch
        self.feature_length = base_ch * 4

        self.net = nn.Sequential(
            nn.Conv2d(input_ch, base_ch, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(base_ch, base_ch * 2, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(base_ch * 2, self.feature_length, kernel_size=3,
                      padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(output_size=(1, 1))
        )
        self.fc = nn.Linear(in_features=self.feature_length,
                            out_features=num_classes)

    def forward(self, input):
        output = self.net(input)
        output = output.view(-1, self.feature_length)
        output = self.fc(output)
        return output


def load_dataset():
    train_dataset = datasets.MNIST(root='./data',
                                   train=True,
                                   transform=transforms.ToTensor(),
                                   download=True)
    test_dataset = datasets.MNIST(root='./data',
                                  train=False,
                                  transform=transforms.ToTensor(),
                                  download=True)
    return train_dataset, test_dataset


def hook_func(module, input, output):
    """
    Hook function of register_forward_hook

    Parameters:
    -----------
    module: module of neural network
    input: input of module
    output: output of module
    """
    image_name = get_image_name_for_hook(module)
    data = output.clone().detach()
    data = data.permute(1, 0, 2, 3)
    vutil.save_image(data, image_name, pad_value=0.5)


def get_image_name_for_hook(module):
    """
    Generate image filename for hook function

    Parameters:
    -----------
    module: module of neural network
    """
    os.makedirs(INSTANCE_FOLDER, exist_ok=True)
    base_name = str(module).split('(')[0]
    index = 0
    image_name = '.'  # '.' is surely exist, to make first loop condition True
    while os.path.exists(image_name):
        index += 1
        image_name = os.path.join(
            INSTANCE_FOLDER, '%s_%d.png' % (base_name, index))
    return image_name


if __name__ == '__main__':
    time_beg = time.time()

    train_dataset, test_dataset = load_dataset()
    train_loader = DataLoader(dataset=train_dataset,
                              batch_size=TRAIN_BATCH_SIZE,
                              shuffle=True)
    test_loader = DataLoader(dataset=test_dataset,
                             batch_size=TEST_BATCH_SIZE,
                             shuffle=False)

    model = Model(input_ch=1, num_classes=10, base_ch=BASE_CHANNEL).cuda()
    torchsummary.summary(
        model, input_size=(INPUT_CHANNEL, INPUT_SIZE, INPUT_SIZE))
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    train_loss = []
    for ep in range(EPOCH):
        # ----------------- train -----------------
        model.train()
        time_beg_epoch = time.time()
        loss_recorder = []
        for data, classes in train_loader:
            data, classes = data.cuda(), classes.cuda()
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, classes)
            loss.backward()
            optimizer.step()

            loss_recorder.append(loss.item())
            time_cost = time.time() - time_beg_epoch
            print('\rEpoch: %d, Loss: %0.4f, Time cost (s): %0.2f' % (
                ep, loss_recorder[-1], time_cost), end='')

        # print train info after one epoch
        train_loss.append(loss_recorder)
        mean_loss_epoch = torch.mean(torch.Tensor(loss_recorder))
        time_cost_epoch = time.time() - time_beg_epoch
        print('\rEpoch: %d, Mean loss: %0.4f, Epoch time cost (s): %0.2f' % (
            ep, mean_loss_epoch.item(), time_cost_epoch), end='')

        # save model
        os.makedirs(MODEL_FOLDER, exist_ok=True)
        model_filename = os.path.join(MODEL_FOLDER, 'epoch_%d.pth' % ep)
        torch.save(model.state_dict(), model_filename)

        # ----------------- test -----------------
        model.eval()
        correct = 0
        total = 0
        for data, classes in test_loader:
            data, classes = data.cuda(), classes.cuda()
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += classes.size(0)
            correct += (predicted == classes).sum().item()
        print(', Test accuracy: %0.4f' % (correct / total))

    print('Total time cost: ', time.time() - time_beg)

    # ----------------- visualization -----------------
    # clear output folder
    if os.path.exists(IMAGE_FOLDER):
        shutil.rmtree(IMAGE_FOLDER)

    model.eval()
    modules_for_plot = (torch.nn.ReLU, torch.nn.Conv2d,
                        torch.nn.MaxPool2d, torch.nn.AdaptiveAvgPool2d)
    for name, module in model.named_modules():
        if isinstance(module, modules_for_plot):
            module.register_forward_hook(hook_func)

    test_loader = DataLoader(dataset=test_dataset,
                             batch_size=1,
                             shuffle=False)
    index = 1
    for data, classes in test_loader:
        INSTANCE_FOLDER = os.path.join(
            IMAGE_FOLDER, '%d-%d' % (index, classes.item()))
        data, classes = data.cuda(), classes.cuda()
        outputs = model(data)

        index += 1
        if index > 20:
            break

程序解释

上述代码简单地训练了一个MNIST分类模型，然后使用最后一个epoch的模型对test数据跑前向，并把特征图保存为图片。

代码中torchsummary仅用来打印模型结构和参数，相关部分可以删掉，不影响运行结果。如果发现缺少该模块，又不想删除的话，安装方式是pip install torchsummary。

训练设置

模型： 三层卷积加一个全连接
Loss：交叉熵
优化器：Adam，学习率=1e-3
训练batch_size：64
（这些都不重要）

重要的内容在visualization部分：

1. test_loader的batch_size设置为1。
2. 下面一段代码用来设置register_forward_hook，只要我们在跑前向之前将register_forward_hook设置好，那么在它就会在前向的时候被调用。

modules_for_plot = (torch.nn.ReLU, torch.nn.Conv2d,
                    torch.nn.MaxPool2d, torch.nn.AdaptiveAvgPool2d)
for name, module in model.named_modules():
    if isinstance(module, modules_for_plot):
        module.register_forward_hook(hook_func)

3. hook_func(module, input, output)需要自己实现。该函数的参数是固定的，不能自己随意加减，这就带来了一个比较麻烦的事情，我们没法把要保存的文件名等信息作为参数传给hook_func，所以这里我使用了一种比较粗暴的方式。文件名的前缀是module的类型名，后面会跟一个序号，序号是同类module在网络中出现的顺序。然而这个序号也没法传给hook_func，所以我直接从硬盘上从序号1开始，通过判断文件的存在性来决定下一个序号，由get_image_name_for_hook函数实现。比较粗暴和无奈的一种方法，如果哪位老兄有更好的方式欢迎留言指教。
4. hook_func中，在保存图片之前，需要将input或output的维度从[1, C, H, W]调整为[C, 1, H, W]
5. save_image()每一行的特征图数量默认是8，可以通过nrow参数进行修改。pad_value是特征图中间填充的间隙的颜色，尽管PyTorch将其定义为整数，但是由于Tensor一般是float型，所以还是输入个[0, 1]之间的浮点数才能真正生效。

结果

看一下三个卷积层的特征图，由于pooling的原因，三个特征图尺寸逐渐减小。

上面图像中我们可以看到一些纯黑的特征图，这些就是所谓的dead neurons，可能称之为dead feature map更合适。下面我们做一下横向比较，即不同输入数据情况下的某一层特征图。通过下图可以看到，不同的输入数据情况下，第二个卷积层的dead feature map位置均相同，这些特征图没有办法提供有效信息，又因它们位置固定，所以其实我们可以将它们对应的卷积核从网络中剔除出去，这样可以起到模型压缩的作用。