Torch 模型 感受野可视化

前言：感受野是卷积神经网络 (CNN) 中一个重要的概念，它表示 CNN 每一层输出的特征图上的像素点在输入图像上映射的区域。感受野的大小和形状直接影响到网络对输入图像的感知范围和精度，进而调整网络结构、卷积核大小和步长等参数，以改善网络的性能。

效果：本文的实验在 torchvision.models 中的 resnet18 上进行，分别绘制了理论感受野、训练前感受野、训练后感受野

开发环境：PyTorch 1.9.0

适用模型：最大池化层使用 nn.MaxPool 而不是 torch.nn.functional.max_pool 的模型

声明：本文所使用代码不开源，觉得本文的思路可行的话，请加 QQ - 1398173074 购买 (￥40，注明来意)

商品仅包含一份 120+ 行的代码。本文所使用的代码基于 torch、matplotlib 以及其它标准库。其中包含一个名为 ReceptiveField 的类，用于绘制图像识别网络的感受野

代码实现

ReceptiveField 提供了以下函数：

• _replace：将 MaxPool (这种求最大值的操作会影响感受野的正确性) 替换为 AvgPool
• __init__：注册前向传播的“挂钩”，用于提取目标层的特征图用于反向传播
• _backward：前向推导图像，利用“挂钩”获取特征图，从特征图中心点反向传播梯度，进行一系列处理后将梯度图转换为感受野图
• theoretical：结合 _backward 函数求解理论感受野，其结果经过 sum、sqrt 之后即为理论感受野的尺寸
• effective：默认情况下结合 _backward 函数求解训练前感受野 (即随机权重的模型)；给定 state_dict 时将加载权重，求解训练后的感受野
• compare：使用 matplotlib 绘制理论感受野、训练前感受野、训练后感受野

class ReceptiveField:
    """ :param model: 需要进行可视化的模型
        :param tar_layer: 感兴趣的层, 其所输出特征图需有 4 个维度 [B, C, H, W]
        :param img_size: 测试时使用的图像尺寸"""

    def make_input(self, n_sample): ...

    def __init__(self,
                 model: nn.Module,
                 tar_layer: Union[int, nn.Module],
                 img_size: Union[int, Tuple[int, int]],
                 use_cuda: bool = False,
                 use_copy: bool = False): ...

    def compare(self, theoretical=True, original=True, state_dict=None, **imshow_kw):
        """ :param theoretical: 是否绘制理论感受野
            :param original: 是否绘制训练前的感受野
            :param state_dict: 模型权值, 如果提供则绘制训练后的感受野"""

    def effective(self, state_dict=None):
        """ :param state_dict: 模型权值, 如果提供则绘制训练后的感受野"""

    def theoretical(self, light=1.):
        """ :param light: 理论感受野的亮度 [0, 1]"""

    def _replace(self, model): ...

    def _backward(self, x): ...

在本文的示例中，对 resnet18 的 layer3 进行了可视化，并计算出理论感受野的尺寸为 211×211

if __name__ == "__main__":
    from torchvision.models import resnet18

    # Step 1: 刚完成初始化的模型, 权重<完全随机>, 表 "训练前"
    m = resnet18()

    # Step 2: 训练完成后的 state_dict, 等待 ReceptiveField 加载
    state_dict = resnet18(pretrained=True).state_dict()

    # Step 3: 绘制感受野 (设置 ReceptiveField 的 use_copy=True, 将创建模型的深拷贝副本)
    with ReceptiveField(m, tar_layer=m.layer3, img_size=256, use_copy=True) as r:
        r.compare(state_dict=state_dict)
        # 理论感受野的尺寸
        s = round(r.theoretical().sum() ** 0.5)
        print(f"Theoretical RF: {s}×{s}")
    plt.show()

    # Step 4: 加载模型的参数
    m.load_state_dict(state_dict)

如果将 resnet18 中的某一个卷积改成空洞卷积，感受野将进一步增大到 243×243

if __name__ == "__main__":
    from torchvision.models import resnet18

    # Step 1: 刚完成初始化的模型, 权重<完全随机>, 表 "训练前"
    m = resnet18()
    print(m)
    m.layer3[1].conv1.dilation = 2
    m.layer3[1].conv1.padding = 2

    # Step 2: 训练完成后的 state_dict, 等待 ReceptiveField 加载
    state_dict = resnet18(pretrained=True).state_dict()

    # Step 3: 绘制感受野 (设置 ReceptiveField 的 use_copy=True, 将创建模型的深拷贝副本)
    with ReceptiveField(m, tar_layer=m.layer3, img_size=256, use_copy=True) as r:
        r.compare(state_dict=state_dict)
        # 理论感受野的尺寸
        s = round(r.theoretical().sum() ** 0.5)
        print(f"Theoretical RF: {s}×{s}")
    plt.show()

    # Step 4: 加载模型的参数
    m.load_state_dict(state_dict)