Day 47 - 注意力热力图 (Attention Heatmap)

一、引言：为什么需要关注模型"看"到了哪里？

在深度学习中，模型往往被视为一个"黑盒"。虽然它能给出很高的分类准确率，但我们很难知道它是基于什么依据做出的判断。

例如，一个识别"狗"的模型，是真正识别出了狗的特征，还是仅仅记住了"草地背景"通常与狗同时出现？

注意力热力图 (Attention Heatmap) 就是一种打开这个黑盒的手段。它将模型对图像各区域的关注程度可视化：

高亮区域 (通常为红色)：表示模型认为该区域对分类决策最重要。
暗淡区域 (通常为蓝色)：表示模型忽略的区域。

通过观察热力图，我们可以：

验证模型逻辑：确认模型是否关注了正确的主体（如狗的头部），而不是背景。
发现数据问题：识别是否存在背景偏差（如所有狼都在雪地里，模型可能其实是在识别雪）。
解释模型决策：向非技术人员直观展示模型的判断依据。

二、核心实现原理

本次笔记的核心是利用 Hook (钩子) 机制 和 通道重要性权重 来生成热力图。

1. 捕获特征图 (Hook 机制)

我们需要获取模型深层（通常是最后一个卷积层）的输出特征图。因为深层特征图包含了最丰富的高级语义信息（如物体的部件、形状）。

PyTorch 提供了 register_forward_hook，可以在模型前向传播时，自动"钩取"中间层的输出。

2. 计算通道权重

特征图通常有多个通道（例如 128 个），每个通道关注不同的特征。我们需要知道哪些通道对当前图像最重要。

类似于 SE (Squeeze-and-Excitation) 模块的思想，我们可以对特征图进行全局平均池化，得到每个通道的平均响应值。响应值越大，说明该通道检测到的特征在当前图像中越显著。

3. 热力图生成与叠加

选出权重最高的几个通道，将其对应的二维特征图提取出来。由于特征图尺寸通常远小于原图（例如 32x32 vs 224x224），我们需要将其上采样 (Resize/Zoom) 到原图大小，并以半透明的方式叠加在原图上显示。

三、代码实现详解

以下是完整的可视化函数代码，包含了从特征提取到绘图的全过程。

复制代码

import torch
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import zoom

def visualize_attention_map(model, test_loader, device, class_names, num_samples=3):
    """
    可视化模型的注意力热力图，展示模型关注的图像区域
    
    参数:
        model: 训练好的 CNN 模型
        test_loader: 测试数据加载器
        device: 计算设备 (CPU/GPU)
        class_names: 类别名称列表
        num_samples: 可视化的样本数量
    """
    model.eval()  # 必须设置为评估模式
    
    with torch.no_grad():
        for i, (images, labels) in enumerate(test_loader):
            if i >= num_samples:
                break
                
            images, labels = images.to(device), labels.to(device)
            
            # -------------------------------------------------------
            # 1. 注册钩子 (Register Hook)
            # -------------------------------------------------------
            activation_maps = []
            
            def hook(module, input, output):
                # 将特征图保存到列表中，注意要转回 CPU
                activation_maps.append(output.cpu())
            
            # 为模型的最后一个卷积层 (这里假设是 conv3) 注册钩子
            # 注意：实际使用时需根据模型结构修改层名称
            hook_handle = model.conv3.register_forward_hook(hook)
            
            # -------------------------------------------------------
            # 2. 前向传播 (Forward Pass)
            # -------------------------------------------------------
            outputs = model(images)
            
            # 务必移除钩子，防止内存泄漏或影响后续操作
            hook_handle.remove()
            
            # 获取预测类别
            _, predicted = torch.max(outputs, 1)
            
            # -------------------------------------------------------
            # 3. 数据预处理与权重计算
            # -------------------------------------------------------
            # 还原原始图像用于显示 (假设做了标准化处理)
            img = images[0].cpu().permute(1, 2, 0).numpy()
            img = img * np.array([0.2023, 0.1994, 0.2010]).reshape(1, 1, 3) + \
                  np.array([0.4914, 0.4822, 0.4465]).reshape(1, 1, 3)
            img = np.clip(img, 0, 1) # 限制值在 [0, 1] 范围
            
            # 获取特征图: [Batch, Channel, Height, Width] -> [Channel, H, W]
            feature_map = activation_maps[0][0] 
            
            # 计算通道注意力权重: 对空间维度 (H, W) 求均值 -> [Channel]
            # 这代表了每个通道在整张图上的平均激活强度
            channel_weights = torch.mean(feature_map, dim=(1, 2))
            
            # 按权重从大到小排序，获取最活跃的通道索引
            sorted_indices = torch.argsort(channel_weights, descending=True)
            
            # -------------------------------------------------------
            # 4. 绘图 (Plotting)
            # -------------------------------------------------------
            fig, axes = plt.subplots(1, 4, figsize=(16, 4))
            
            # 子图1: 原始图像与预测结果
            axes[0].imshow(img)
            axes[0].set_title(f'Original Image\\nTrue: {class_names[labels[0]]}\\nPred: {class_names[predicted[0]]}')
            axes[0].axis('off')
            
            # 子图2-4: 显示前3个最活跃通道的热力图
            for j in range(3):
                channel_idx = sorted_indices[j]
                
                # 获取该通道的二维特征图
                channel_map = feature_map[channel_idx].numpy()
                
                # 归一化到 [0, 1]，保证热力图颜色分布正常
                channel_map = (channel_map - channel_map.min()) / (channel_map.max() - channel_map.min() + 1e-8)
                
                # 上采样: 将小的特征图 (如 32x32) 放大到原图尺寸 (如 32x32 -> 原图大小)
                # 注意：这里 zoom 的系数是根据特征图尺寸和目标尺寸计算的
                # 若原图很大，这里需要调整缩放比例
                heatmap = zoom(channel_map, (img.shape[0]/feature_map.shape[1], img.shape[1]/feature_map.shape[2]))
                
                # 叠加显示
                axes[j+1].imshow(img)
                # alpha=0.5 设置半透明，cmap='jet' 使用经典的蓝-红热力图配色
                axes[j+1].imshow(heatmap, alpha=0.5, cmap='jet')
                axes[j+1].set_title(f'Attention Heatmap\\nChannel {channel_idx}')
                axes[j+1].axis('off')
            
            plt.tight_layout()
            plt.show()

四、结果解读与分析

当你运行上述代码后，会看到一系列图片，每一组包含一张原图和三张热力图。

1. 热力图颜色含义

红色/深红：高响应区域。这是模型"最在意"的地方。
蓝色/深蓝：低响应区域。这是模型认为无关紧要的背景。

2. 多通道的互补性

你会发现，不同的通道关注点往往不同：

通道 A 可能聚焦于物体的轮廓边缘。
通道 B 可能聚焦于特定的纹理（如斑马的条纹、鸟的羽毛）。
通道 C 可能聚焦于特定部位（如眼睛、车轮）。

这就是神经网络的"分工合作"。最终的分类结果是所有这些通道特征综合作用的产物。

3. 如何判断模型好坏？

好模型：热力图的红色区域紧密地覆盖在目标物体上。例如，识别"猫"时，红色集中在猫的身体上，背景基本为蓝色。
坏模型/过拟合：热力图散乱，或者错误地聚焦在背景上。例如，红色区域出现在天空或草地上，说明模型学到了错误的特征关联。

五、总结

通过可视化注意力热力图，我们将抽象的神经网络特征转化为了人类可理解的视觉信息。

这不仅增强了我们对模型的信任度，也为后续的模型优化（如针对错误关注区域进行数据增强）提供了明确的方向。

核心技术点在于利用 hook 获取中间层输出，并利用全局平均池化计算通道的重要性权重。