【PyTorch】获取特征图并可视化

0. 前言

都说神经网络是个黑盒子，其内部具体的工作原理也模糊不清，但是，对卷积神经网络输出的特征图进行可视化能够大概地知悉网络每个阶段的作用，能够为网络设计提供一些思路。在此记录一下工作中常用的特征图提取和可视化方法。

1. 特征图提取

博主常用的特征图提取方法有3种：

• 在forward函数中直接return，优点是快速，简单，缺点是需要修改网络部分的代码，并且无法获取nn.Sequential子模块的特征图。
• 使用register_forward_hook，可以获取指定的一层或多层输出的特征图。这篇文章中有相关的代码。
• 使用torchvision.models.feature_extraction.create_feature_extractor，这是目前我认为最方便的特征图提取方法，可以指定任意一层或多层。

方法3的官方示例代码：

>>> from torchvision.models.feature_extraction import get_graph_node_names, create_feature_extractor
>>> # Feature extraction with resnet
>>> model = torchvision.models.resnet18()
>>> # 获取节点名称
>>> nodes, _ = get_graph_node_names(model)
>>> # extract layer1 and layer3, giving as names `feat1` and feat2`
>>> model = create_feature_extractor(
>>>     model, {'layer1': 'feat1', 'layer3': 'feat2'})
>>> out = model(torch.rand(1, 3, 224, 224))
>>> print([(k, v.shape) for k, v in out.items()])
>>>     [('feat1', torch.Size([1, 64, 56, 56])),
>>>      ('feat2', torch.Size([1, 256, 14, 14]))]

2. 特征图可视化

调用matplotlib.pyplot库，通过plt.imshow()函数逐通道显示特征图或者合并通道显示。

参考

https://pytorch.org/vision/stable/feature_extraction.html