SlowFast/mmaction2 中的 GradCAM 工具

0. 前言

• 最近一直在mmaction2上进行二次开发，有一个需求，想实现一个gradcam工具来分析模型训练结果。
  • 目前实现在行为识别模型上的功能实现，已提交PR
  • 未来希望能在时空行为检测上实现功能。为了给到时候的自己一些帮助，记录一下整体过程。
• 为了实现上面这个功能，想找点参考资料，正好SlowFast有这个功能（虽然文档很不完善，但总归有整个流程），所以现在系统学习一下这个功能。
• 所以本文主要内容：
  • 介绍SlowFast中的Gradcam工具的实现与使用。
  • 介绍在 mmaction2 中复现的过程。

1. SlowFast 中的 GradCAM

• 可视化结果可以参考 SlowFast Demo

1.1. GradCAM的功能

• 输入：一个训练好的模型和一组输入数据（一般是一张或一组图像，可以指定label，也可以不指定）。
• 输出：每张图像每个像素位置的重要性权重，每张图片对应一个重要性权重矩阵。
  • 所谓重要性权重，指的是每个像素对结果的影响大小，权重越大影响越大。
  • 注意，重要性权重矩阵其实是一个灰度图（即shape为[height, width, 1]），可通过matplotlib将灰度图转换为热力图。
  • 将热力图与原始图片叠加，就能得到SlowFast Demo中的效果。

1.2. GradCAM 的原理

• 详情请参考原论文，这里只说一些与实现相关的内容。

• 重要性权重矩阵的本质是前向特征图加权和+resize的结果。

• 所谓前向特征图加权和就包括两个部分

  • 这一部分这里只是简单描述，后面在 mmaction2 复现的过程中会详细介绍。
  • 前向特征图：需要在Gradcam工具中选中网络中的某一层，通过前向传播得到这一层的输出。提取层的shape为一般为 [batch_size, num_segments, channels, height, width]
  • 特征图每一层的权重：权重是选中梯度层反向传播得到的梯度矩阵，进行global avg pool的结果，结果一般为 [batch_size, num_segments, channels]
  • 通过加权可以得到 [batch_size, num_segments, feature_height, feature_width]

• resize的功能

  • 前向特征图加权和的shape与原始图像不同，resize的目标是将上一部得到的前向特征图转换到需要的尺寸。
  • 需要注意的是，resize的操作不一定只有高度与宽度，可能还有时间纬度。

• GradCAM主要超参数就是选择网络中的特定层，作为前向特征图加权和的输入。

1.3. SlowFast 中的具体实现

• 运行阶段：SlowFast中，热力图生成相关函数有独立于训练、验证、测试存在，有一个叫visualization.py的脚本。
• GradCAM的实现其实就只有一个类
• 主要需要实现的功能有：
  • 构建输入数据（一般是经过crop+sampling+norm）
  • 根据输入数据与训练好的模型，获取指定特征图的前向结果以及反向传播梯度值。
  • 根据前向以及反向结果，构建重要性权重矩阵。
  • 将输入数据（需renorm操作）与特征图矩阵融合，得到最终展示结果。
• 具体实现细节会在 mmaction2 中介绍一些。

2. mmaction2 复现 GradCAM 功能

• 已提交PR

2.1. mmaction2 中的 shape

• 数据shape以及各种transpose操作在实现GradCAM时非常重要，趁这个机会总结一下。
• 数据预处理中的shape与transpose
  • 第一步，数据生成，一般是在mmaction.datasets.pipelines.loading.py中的各种decode方法。
    • 这一步的结果是 list，其中每个元素代表一帧图像，图像由ndarray表示。
    • 视频是DecordDecode/PyAVDecord/OpenCVDecode方法，通过不同库从视频文件中读取特定帧，组成列表。
    • 图像帧是 RawFrameDecode 方法，通过帧文件夹路径、图像文件名、帧编号来构建帧列表。
  • 第二步，数据预处理一般包括各种crop操作，有两类操作：
    • 不改变shape的crop操作：MultiScaleCrop/RandomResizedCrop/RandomCrop/CenterCrop
    • 改变shape的crop操作：ThreeCrop/TenCrop/MultiGroupCrop。
      • 本质就是一张图片经过crop变为多张图片。
      • 值得一提的是，shape可reshape为[num_crops, num_imgs, height, width, 3]，即list是由num_crops组数据组成的，每组图片都是相同crop bbox在不同图片上的切片。
  • 第三步，一般会进行norm操作，不会对shape有任何影响。
  • 第四步，会进行FormatShape操作，这个操作本质就是shape以及transpose操作。分为两大类
    • NCHW模式：用于2D模型，如TSN/TSM/TIN，模型的输入数据shape为[batch_size, num_segments, 3, img_height, img_width]。
    • NCTHW模式：用于3D模型，如I3D/SlowFast/C3D等。
      • 输入shape中多了 [batch_size, num_cropsxnum_clips, 3, clip_len, img_height, img_width]
      • 其中，num_clips, clip_len都是之前提取帧编号SampleFrames的输出结果。
      • num_crops是通过resize纬度为-1生成的，跟上一步的crop类别有关。
  • 之后的操作与shape/transpose也没太多关系。
• 模型中的shape与transpose
  • 模型分为2D模型与3D模型，每类模型都分为三个部分，backbone/neck/head/average_clip。
  • 2D模型，即Recognizer2D，如TSN/TSM/TIN。
    • 模型输入五纬数据，即[batch_size, num_segments, 3, img_height, img_width]
    • backbone的输入为四位数据，即[batch_size*num_segments, 3, img_height, img_width]，中间特征图的尺寸也都是四位的，且第一纬不会变化，就是普通的2D卷积。
    • neck用得比较少，就TPN里用了，但我也没仔细看过源码，这里就不介绍了。
    • head的输入也是四位数据，即[batch_size*num_segments, C, H, W]，输出 [batch_size, num_classes]
  • 3D模型，即Recognizer3D，如I3D/CSN/X3D/SlowOnly等。
    • 输入六纬数据，即[batch_size, num_crops*num_clips, 3, clip_len, img_height, img_width]
    • backbone的输入为五维数据，即[batch_size*num_crops*num_clips, 3, clip_len, img_height, img_width]，中间特征图也是5维，第一纬不会变化。
    • neck用得比较少，就TPN里用了，但我也没仔细看过源码，这里就不介绍了。
    • head的输入是五维数据，[batch_size*num_crops*num_clips, C, T, H, W]，输出[batch_size*num_crops*num_clips, num_classes]

2.2. GradCAM 的具体实现

• 参考SlowFast中的实现，复现GradCAM也可以分为三步：
  • 第一步：获取指定layer的正向与反向结果。
  • 第二步：根据正向、反向结果，构建重要性权重矩阵。
  • 第三步：融合重要性权重矩阵与输入图像，得到最终结果。
• 第一步：获取指定layer的正向与反向结果。
  • 通过 nn.Module的hook功能实现，即注册 register_backward_hook与register_forward_hook。
  • 前者在每次调用完forward方法后调用，后者在每次调用完backward方法后调用。
  • 更多内容可以参考文档
• 第二步：根据正向、反向结果，构建重要性权重矩阵。
  • 目标：根据第一步输出，即前向结果activations与反向结果gradients的，为每张图片输出重要性权重矩阵，即一个灰度图。
  • 实现过程难度低，但很麻烦，主要就是各种shape，各种transpose。
  • 有几个需要注意的点：
    • mmaction2中的2D模型与3D模型输入数据的shape不相同，需要分别处理。
    • 我们会对重要性权重矩阵进行norm操作
      • $$normedmap=\frac{map-min(map)}{max(map)-min(map)+\Delta }$$
      • 其中 $\Delta$是极小值，防止分母为0，SlowFast里设置这个数值为1e-6。
      • 但在实际使用中会出现max(map) - min(map)非常小的情况（比如我就碰到过这个数值差不多是1e-18），即$\Delta >>max(map)-min(map)$，这就会导致normedmap整体数值非常接近，热力图基本上都是同一个颜色。
      • 在实现的时候，将这个数值减小到1e-20，虽然暂时解决了上面的问题，但我在想，梯度都这么小了，能体现出热力图本来的效果吗？
• 第三步：融合重要性权重矩阵与输入图像，得到最终结果。
  • 重要性权重矩阵是灰度图，为了变为热力图的形式，需要进行色彩转换。
    • matplotlib就提供了这方面的功能，本质就是使用plt.get_cmap(colormap)，生成的彩色图为RGB格式，像素值在[0, 1]之间。
    • 更多信息可以参考matplotlib colormap文档
  • 模型输入数据中，图像是经过norm操作的，即已经减去均值、除以方差。
    • 为了融合图像，那输入数据与热力图的形式必须相同，即我们需要将输入数据也转换成上面提到的：RGB格式，像素值在[0, 1]之间。
    • 需要注意的是，mmaction2中原始图像像素没有先转换到[0, 1]之间，而是直接在[0, 255]这个范围上进行norm操作的，所以在renorm之后还需要对所有像素值除以255。
  • 图像融合其实用的是 alpha blending，看起来很高端，其实就是两个图像的加权和，权重就是$\alpha$与$1-\alpha$。

2.3. 其他

• 记录一下心情。
• mmaction2真是高标准严要求，code review在线挨打，主要在于文档编写不规范，还有各种细节（修改意见至少包括6处标点符号）。
• 写功能大概花了1-2天全天，写文档+单元测试又花了1天……如果自己用，其实我只用其中几个模型，就不用写测试，到时候随便调一调代码就好。给别人用，就需要把所有模型都试一遍。只能说，我对mmaction2的源码理解更深了……
• 虽然一直在挨打，但还是有一些收获的，相信如果还要提交PR，不会像这次这么惨了。
• GradCAM并没有那么好，要好看的效果也需要稍微调一调：
  • 各种layer要都试一试
  • delta也许过大
  • 输入图像的size可能需要调一调