深度学习指标可视化案例

时间:2025-02-03 07:01:12
代码案例:
import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 加载图像
img = Image.open('path_to_image.jpg')
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
img_tensor = transform(img).unsqueeze(0)

# 获取模型的最后一个卷积层的输出
def get_last_conv_layer(model):
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            last_conv_layer = module
    return last_conv_layer

last_conv_layer = get_last_conv_layer(model)

# 前向传播
outputs = []
def hook(module, input, output):
    outputs.append(output)
handle = last_conv_layer.register_forward_hook(hook)

with torch.no_grad():
    output = model(img_tensor)
    handle.remove()

# 获取特征图和类别权重
feature_map = outputs[0].squeeze(0).cpu().numpy()
weights = model.fc.weight.data.cpu().numpy()

# 计算 Grad-CAM
class_idx = torch.argmax(output).item()
cam = np.zeros(feature_map.shape[1:], dtype=np.float32)
for i, w in enumerate(weights[class_idx]):
    cam += w * feature_map[i, :, :]

cam = cv2.resize(cam, (224, 224))
cam = np.maximum(cam, 0)
cam = cam / np.max(cam)

# 可视化
img = cv2.imread('path_to_image.jpg')
img = cv2.resize(img, (224, 224))
heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
superimposed_img = heatmap * 0.4 + img * 0.6

plt.imshow(superimposed_img.astype(np.uint8))
plt.axis('off')
plt.show()
此代码通过 Grad-CAM 生成热力图,可视化模型对输入图像的注意力区域。