深度学习中获得中间的特征图/low level feature map/high level feature map/特征图可视化

x.1 low/high level feature map

low level feature指的是角点，边缘信息，是浅层卷积核学习的，即一开始的卷积核学习的。high level feature指的是纹理等抽象信息，是深层卷积核学习的，即越后面的卷积核学习的越抽象。

CNN具有比较强的可解释性，为什么呢？有个很大的原因就在于它的中间特征矩阵是可以可视化的，像Attention的中间特征矩阵就不好可视化。

下面要讲的获得CNN中间特征图的方法大致分为两种，一种是在model中直接定义一个list用于保存中间特征矩阵，然后返回，简称为“重写model”方法；一种是使用pytorch自定义的注册hook进行输出中间特征图，简称为“注册hook”方法。前者适用于网络较为复杂，如采用多层递归等等，但是需要重写forward函数，用于保存你需要的特征矩阵；后者简洁，且不需要重写模型中的forward函数。

x.2 重写model方法

重写model方法中，即使用一个list变量存储你需要特征矩阵并在forward函数中，最终return你的特征矩阵。

import torch.nn as nn
import torch


class AlexNet(nn.Module):
    def __init__(self, num_classes=1000, init_weights=False):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  # input[3, 224, 224]  output[48, 55, 55]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[48, 27, 27]
            nn.Conv2d(48, 128, kernel_size=5, padding=2),           # output[128, 27, 27]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 13, 13]
            nn.Conv2d(128, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 192, kernel_size=3, padding=1),          # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 128, kernel_size=3, padding=1),          # output[128, 13, 13]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),                  # output[128, 6, 6]
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(128 * 6 * 6, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        outputs = []
        for name, module in self.features.named_children():
            x = module(x)
            if name in ["0", "3", "6"]:
                outputs.append(x)

        return outputs

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

x.3 注册hook方法

x.3.1 使用hook方法输出特征图大小的例子

在PyTorch中，可以使用hook函数来获取神经网络中某一层的输出（也就是所谓的“特征图”或“feature map”），以便进行后续的分析或处理。下面是一个示例代码，展示了如何使用hook函数获取某一层的输出：

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.conv3(x)
        x = self.relu(x)
        return x

# 定义一个hook函数，用于获取某一层的输出
def get_feature_map(module, input, output):
    print('Feature map size:', output.size())

# 创建一个模型实例，并注册hook函数
model = MyModel()
hook_handle = model.conv2.register_forward_hook(get_feature_map)

# 创建一个输入张量，并通过模型进行前向传播
x = torch.randn((1, 3, 32, 32))
output = model(x)

# 打印输出结果
print('Model output size:', output.size())

# 移除hook函数
hook_handle.remove()

在上述示例中，我们定义了一个MyModel类，该类包含了三个卷积层和一个ReLU激活函数。然后，我们使用register_forward_hook方法注册了一个hook函数，用于获取第二个卷积层的输出。最后，我们传入一个输入张量x，并通过模型进行前向传播，同时在hook函数中获取了第二个卷积层的输出大小。

需要注意的是，在hook函数中我们可以获取到模块的输入、输出以及模块自身，这些参数都可以在后续的分析或处理中使用。在本例中，我们只是简单地打印了特征图的大小，但实际上可以根据需要对特征图进行进一步的操作，比如可视化或保存到文件中等。

x.3.2 将hook获得的特征图可视化

进行可视化即更改hook函数。

对于上述示例中获取到的特征图，可以使用Matplotlib库将其可视化，或使用PIL库将其保存到文件中。下面是示例代码：

# 使用Matplotlib将特征图可视化
import matplotlib.pyplot as plt
import numpy as np

def visualize_feature_map(module, input, output):
    # 将Tensor转换为NumPy数组，并取第一个样本
    feature_map = output.detach().numpy()[0]
    # 获取特征图的通道数
    num_channels = feature_map.shape[0]
    # 创建一个网格图，每行显示8个特征图
    fig, axs = plt.subplots(nrows=int(np.ceil(num_channels/8)), ncols=8, figsize=(12, 6))
    for i in range(num_channels):
        # 获取第i个特征图，并将其归一化到0-1之间
        channel_map = feature_map[i, ...]
        channel_map -= np.min(channel_map)
        channel_map /= np.max(channel_map)
        # 在网格图中显示第i个特征图
        row = i // 8
        col = i % 8
        axs[row][col].imshow(channel_map, cmap='gray')
        axs[row][col].set_xticks([])
        axs[row][col].set_yticks([])
    plt.show()

# 注册hook函数并进行前向传播
hook_handle = model.conv2.register_forward_hook(visualize_feature_map)
output = model(x)

# 移除hook函数
hook_handle.remove()

在上述代码中，我们定义了一个新的hook函数visualize_feature_map，用于将特征图可视化。在该函数中，我们首先将输出张量转换为NumPy数组，并取出第一个样本的特征图。然后，我们使用Matplotlib库创建一个网格图，并将每个特征图归一化到0-1之间后在网格图中显示。最后，我们通过调用plt.show()方法显示可视化结果。

x.3.3 将hook获得的特征图保存

除了将特征图可视化外，还可以使用PIL库将其保存到文件中，如下所示：

# 使用PIL将特征图保存到文件中
from PIL import Image

def save_feature_map(module, input, output):
    # 将Tensor转换为PIL图像，并保存到文件中
    feature_map = output.detach().numpy()[0]
    for i in range(feature_map.shape[0]):
        channel_map = feature_map[i, ...]
        channel_map -= np.min(channel_map)
        channel_map /= np.max(channel_map)
        channel_map = (channel_map * 255).astype(np.uint8)
        image = Image.fromarray(channel_map)
        image.save(f'channel_{i}.png')

# 注册hook函数并进行前向传播
hook_handle = model.conv2.register_forward_hook(save_feature_map)
output = model(x)

# 移除hook函数
hook_handle.remove()

在上述代码中，我们定义了一个新的hook函数save_feature_map，用于将特征图保存到文件中。在该函数中，我们首先将输出张量转换为NumPy数组，并对每个特征图进行归一化和类型转换，然后使用PIL库创建一个图像对象，并将其保存。

x.3.4 torch.nn.Sequential的特征图可视化

如果网络结构是使用nn.Sequential进行书写的，那么可以通过指定Sequential中子模块的名称或者索引来获取特定层的特征矩阵。

例如，假设您的网络结构如下：

import torch.nn as nn

model = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(256 * 4 * 4, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10)
)

要获取第2个卷积层（即nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)）的输出特征矩阵，可以通过以下代码实现：

hook_handle = model[3].register_forward_hook(hook_fn)

其中，model[3]表示获取Sequential中的第4个子模块（由于Python中的索引从0开始计数，所以model[3]实际上是第4个子模块），并注册hook_fn作为其forward hook函数。在这个hook函数中，您可以像之前的示例一样获取和处理特征矩阵。

如果您想通过子模块名称来获取特定层的特征矩阵，则可以使用nn.Sequential.named_modules()方法。例如，要获取第2个卷积层的输出特征矩阵，可以使用以下代码：

for name, module in model.named_modules():
    if name == '3':
        hook_handle = module.register_forward_hook(hook_fn)

其中，name表示子模块的名称（由于Sequential中子模块没有名称，因此默认情况下，它们的名称为其在Sequential中的索引），module表示子模块本身。在这个循环中，我们检查每个子模块的名称，找到名称为’3’的子模块（即第2个卷积层），并注册hook_fn作为其forward hook函数。注意，named_modules()方法会递归地遍历Sequential中的所有子模块，因此如果您的网络结构比较复杂，可能需要在循环中添加额外的逻辑来确保正确地获取到特定层。

x.3.5 嵌套模型的特征图可视化

如果模型中嵌套了一个小模型，也是继承自nn.Module的，要获取该小模型中某一特定层的输出特征矩阵，可以使用类似于前面的示例的方法。假设您的模型结构如下：

import torch.nn as nn

class SmallModel(nn.Module):
    def __init__(self):
        super(SmallModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.maxpool(x)
        return x


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.small_model = SmallModel()
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(256 * 4 * 4, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = self.small_model(x)
        x = self.conv3(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

要获取SmallModel中的第2个卷积层（即nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)）的输出特征矩阵，可以使用以下代码：

hook_handle = model.small_model.conv2.register_forward_hook(hook_fn)

其中，model.small_model.conv2表示获取SmallModel中的conv2子模块（即第2个卷积层），并注册hook_fn作为其forward hook函数。在这个hook函数中，您可以像之前的示例一样获取和处理特征矩阵。

需要注意的是，您需要在SmallModel的forward方法中将每个子模块逐个调用，以便正确计算输出特征矩阵。在这个例子中，由于SmallModel的forward方法中已经将conv1、relu、conv2和maxpool子模块都调用了一遍，因此在获取SmallModel中的第2个卷积层的输出特征矩阵时，我们只需要关注SmallModel中的conv2子模块。

x.3.6 递归模型的特征图可视化

如果模型是递归定义的，可以使用递归函数来访问特定层的输出特征矩阵。在递归函数中，您需要检查当前层是否是目标层，如果是，则注册hook并获取其输出特征矩阵；如果不是，则继续递归进入下一层。以下是一个示例：

import torch.nn as nn

class RecursiveModel(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(RecursiveModel, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.recursive = None
        if out_channels > 1:
            self.recursive = RecursiveModel(out_channels, out_channels // 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        if self.recursive:
            x = self.recursive(x)
        x = self.conv1(x)
        x = self.relu(x)
        return x

def hook_fn(module, input, output):
    # process the output feature map
    print(output.shape)

def recursive_hook(model, layer_name, current_depth, target_depth):
    for name, module in model._modules.items():
        if isinstance(module, nn.Module):
            if current_depth == target_depth and name == layer_name:
                # if this is the target layer, register the hook
                hook_handle = module.register_forward_hook(hook_fn)
                hooks.append(hook_handle)
            elif current_depth < target_depth:
                # if we haven't reached the target depth, keep going recursively
                recursive_hook(module, layer_name, current_depth + 1, target_depth)

model = RecursiveModel(3, 16)
hooks = []
recursive_hook(model, "conv1", 1, 2)

在这个例子中，我们定义了一个递归模型RecursiveModel，它包含一个卷积层和一个递归子模块。在递归函数recursive_hook中，我们遍历模型中的所有子模块，并检查当前深度是否等于目标深度，以及当前子模块是否是目标层。如果当前深度等于目标深度且当前子模块是目标层，我们就注册forward hook并将其句柄添加到hooks列表中。如果当前深度小于目标深度，我们就递归进入下一层。

在本例中，我们希望获取第2层递归中的第1个卷积层（即模型中的第二个卷积层）。因此，我们在recursive_hook中设置目标深度为2，目标层名称为"conv1"。当我们注册forward hook时，它将捕获目标层的输出特征矩阵。

所以尽量使用顺序的，多嵌套的代码来实现你的模型。