Pytorch可视化特征图

小灰灰 2022-12-20 02:43 182阅读 0赞

### 文章目录 ###

*  为什么要可视化特征图？
 *  PyTorch中与可视化特征图相关的函数/接口
 *   *  torch.nn.Module.register\_forward\_hook
     *  torchvision.utils.make\_grid, torchvision.utils.save\_image
 *  示例：MNIST
 *   *  程序解释
     *  结果

# 为什么要可视化特征图？ #

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

比如其中的代表作之一，2013年Zeiler和Fergus发表的《[Visualizing and Understanding Convolutional Networks][]》  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc_size_16_color_FFFFFF_t_70_pic_center]  
早期LeCun 1998年的文章《[Gradient-Based Learning Applied to Document Recognition][]》中的一张图也非常精彩，个人觉得比Zeiler 2013年的文章更能给人以启发。从下图的`F6`特征，我们可以清楚地看到原始书写差异非常大的图片如何在深层特征中体现出其不变性。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc_size_16_color_FFFFFF_t_70_pic_center 1]

接下来我们看看如何在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)

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

## 结果 ##

看一下三个卷积层的特征图，由于pooling的原因，三个特征图尺寸逐渐减小。  
![在这里插入图片描述][20201111230850662.png_pic_left]  
![在这里插入图片描述][20201111230857765.png_pic_left]  
![在这里插入图片描述][20201111230903945.png_pic_left]

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

[Visualizing and Understanding Convolutional Networks]: https://arxiv.org/abs/1311.2901
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc_size_16_color_FFFFFF_t_70_pic_center]: /images/20221120/ad40a16470754b648edee18da4a33cfa.png
[Gradient-Based Learning Applied to Document Recognition]: http://yann.lecun.com/exdb/publis/pdf/lecun-98.pdf
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc_size_16_color_FFFFFF_t_70_pic_center 1]: https://img-blog.csdnimg.cn/2020111023282996.PNG?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc=,size_16,color_FFFFFF,t_70#pic_center
[20201111230850662.png_pic_left]: https://img-blog.csdnimg.cn/20201111230850662.png#pic_left
[20201111230857765.png_pic_left]: https://img-blog.csdnimg.cn/20201111230857765.png#pic_left
[20201111230903945.png_pic_left]: https://img-blog.csdnimg.cn/20201111230903945.png#pic_left
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc_size_16_color_FFFFFF_t_70_pic_left]: https://img-blog.csdnimg.cn/20201111231645379.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2JieTE5ODc=,size_16,color_FFFFFF,t_70#pic_left