Python打卡 DAY 47

昨天代码中注意力热图的部分顺移至今天

知识点回顾：

热力图

作业：对比不同卷积层热图可视化的结果

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
 
# 设置替代中文字体（适用于Linux）
plt.rcParams["font.family"] = ["WenQuanYi Micro Hei", "sans-serif"]
plt.rcParams['axes.unicode_minus'] = False
 
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CNN()
model.load_state_dict(torch.load('cifar10_cnn_model.pth', map_location=device))
model.to(device)
 
# 4. 将模型设置为评估模式（如果是用于推理）
model.eval()
 
print("模型已成功加载！")
 
class_names = ['飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '卡车']
 
# 可视化空间注意力热力图（显示模型关注的图像区域）
def visualize_attention_map(model, test_loader, device, class_names, num_samples=3):
    """可视化模型的注意力热力图，展示模型关注的图像区域"""
    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)
            
            # 创建一个钩子，捕获中间特征图
            activation_maps = []
            
            def hook(module, input, output):
                activation_maps.append(output.cpu())
            
            # 为最后一个卷积层注册钩子（获取特征图）
            hook_handle = model.conv3.register_forward_hook(hook)
            
            # 前向传播，触发钩子
            outputs = model(images)
            
            # 移除钩子
            hook_handle.remove()
            
            # 获取预测结果
            _, predicted = torch.max(outputs, 1)
            
            # 获取原始图像
            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)
            
            # 获取激活图（最后一个卷积层的输出）
            feature_map = activation_maps[0][0].cpu()  # 取第一个样本
            
            # 计算通道注意力权重（使用SE模块的全局平均池化）
            channel_weights = torch.mean(feature_map, dim=(1, 2))  # [C]
            
            # 按权重对通道排序
            sorted_indices = torch.argsort(channel_weights, descending=True)
            
            # 创建子图
            fig, axes = plt.subplots(1, 4, figsize=(16, 4))
            
            # 显示原始图像
            axes[0].imshow(img)
            axes[0].set_title(f'原始图像\n真实: {class_names[labels[0]]}\n预测: {class_names[predicted[0]]}')
            axes[0].axis('off')
            
            # 显示前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)
                
                # 调整热力图大小以匹配原始图像
                from scipy.ndimage import zoom
                heatmap = zoom(channel_map, (32/feature_map.shape[1], 32/feature_map.shape[2]))
                
                # 显示热力图
                axes[j+1].imshow(img)
                axes[j+1].imshow(heatmap, alpha=0.5, cmap='jet')
                axes[j+1].set_title(f'注意力热力图 - 通道 {channel_idx}')
                axes[j+1].axis('off')
            
            plt.tight_layout()
            plt.show()
 
# 调用可视化函数
visualize_attention_map(model, test_loader, device, class_names, num_samples=3)

@浙大疏锦行