如何在PyTorch自定义脑肿瘤检测CNN模型上实现Grad-CAM？

问题描述 #

我正在开展脑肿瘤检测项目，使用了自定义的CNN_TUMOR模型。希望为该模型添加Grad-CAM可视化功能，但目前多数Grad-CAM实现基于TensorFlow，PyTorch版本的实现多针对预训练模型，而非自定义模型。现请求指导如何在已训练好的该自定义PyTorch模型上实现Grad-CAM。

自定义模型与训练代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
import copy
from tqdm import tqdm

# Define Architecture For CNN_TUMOR Model
class CNN_TUMOR(nn.Module):
    # Network Initialisation
    def __init__(self, params):
        super(CNN_TUMOR, self).__init__()
        Cin,Hin,Win = params["shape_in"]
        init_f = params["initial_filters"]
        num_fc1 = params["num_fc1"]
        num_classes = params["num_classes"]
        self.dropout_rate = params["dropout_rate"]

        # Convolution Layers
        self.conv1 = nn.Conv2d(Cin, init_f, kernel_size=3)
        h,w=findConv2dOutShape(Hin,Win,self.conv1)
        self.conv2 = nn.Conv2d(init_f, 2*init_f, kernel_size=3)
        h,w=findConv2dOutShape(h,w,self.conv2)
        self.conv3 = nn.Conv2d(2*init_f, 4*init_f, kernel_size=3)
        h,w=findConv2dOutShape(h,w,self.conv3)
        self.conv4 = nn.Conv2d(4*init_f, 8*init_f, kernel_size=3)
        h,w=findConv2dOutShape(h,w,self.conv4)

        # compute the flatten size
        self.num_flatten=h*w*8*init_f
        self.fc1 = nn.Linear(self.num_flatten, num_fc1)
        self.fc2 = nn.Linear(num_fc1, num_classes)

    def forward(self,X):
        # Convolution & Pool Layers
        X = F.relu(self.conv1(X));
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv2(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv3(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv4(X))
        X = F.max_pool2d(X, 2, 2)
        X = X.view(-1, self.num_flatten)
        X = F.relu(self.fc1(X))
        X = F.dropout(X, self.dropout_rate)
        X = self.fc2(X)
        return F.log_softmax(X, dim=1)

def findConv2dOutShape(Hin, Win, conv, pool=2):
    kernel_size = conv.kernel_size
    stride = conv.stride
    padding = conv.padding
    dilation = conv.dilation

    Hout = (Hin + 2*padding[0] - dilation[0]*(kernel_size[0]-1) - 1) // stride[0] + 1
    Wout = (Win + 2*padding[1] - dilation[1]*(kernel_size[1]-1) - 1) // stride[1] + 1
    Hout = Hout // pool
    Wout = Wout // pool
    return Hout, Wout

params_model={
        "shape_in": (3,256,256),
        "initial_filters": 8,
        "num_fc1": 100,
        "dropout_rate": 0.25,
        "num_classes": 2}

loss_func = nn.NLLLoss(reduction="sum")

# 假设cnn_model已初始化
cnn_model = CNN_TUMOR(params_model)
opt = optim.Adam(cnn_model.parameters(), lr=3e-4)
lr_scheduler = ReduceLROnPlateau(opt, mode='min',factor=0.5, patience=20,verbose=1)

# 训练相关函数
def get_lr(opt):
    for param_group in opt.param_groups:
        return param_group['lr']

def loss_batch(loss_func, output, target, opt=None):
    loss = loss_func(output, target)
    pred = output.argmax(dim=1, keepdim=True)
    metric_b=pred.eq(target.view_as(pred)).sum().item()

    if opt is not None:
        opt.zero_grad()
        loss.backward()
        opt.step()

    return loss.item(), metric_b

def loss_epoch(model,loss_func,dataset_dl,opt=None):
    run_loss=0.0
    t_metric=0.0
    len_data=len(dataset_dl.dataset)

    for xb, yb in dataset_dl:
        xb=xb.to(device)
        yb=yb.to(device)
        output=model(xb)
        loss_b,metric_b=loss_batch(loss_func, output, yb, opt)
        run_loss+=loss_b

        if metric_b is not None:
            t_metric+=metric_b

    loss=run_loss/float(len_data)
    metric=t_metric/float(len_data)
    return loss, metric

def Train_Val(model, params,verbose=False):
    epochs=params["epochs"]
    loss_func=params["f_loss"]
    opt=params["optimiser"]
    train_dl=params["train"]
    val_dl=params["val"]
    lr_scheduler=params["lr_change"]
    weight_path=params["weight_path"]

    loss_history={"train": [],"val": []}
    metric_history={"train": [],"val": []}
    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss=float('inf')

    for epoch in tqdm(range(epochs)):
        current_lr=get_lr(opt)
        if(verbose):
            print('Epoch {}/{}, current lr={}'.format(epoch, epochs - 1, current_lr))

        model.train()
        train_loss, train_metric = loss_epoch(model,loss_func,train_dl,opt)
        loss_history["train"].append(train_loss)
        metric_history["train"].append(train_metric)

        model.eval()
        with torch.no_grad():
            val_loss, val_metric = loss_epoch(model,loss_func,val_dl)

        if(val_loss < best_loss):
            best_loss = val_loss
            best_model_wts = copy.deepcopy(model.state_dict())
            torch.save(model.state_dict(), weight_path)
            if(verbose):
                print("Copied best model weights!")

        loss_history["val"].append(val_loss)
        metric_history["val"].append(val_metric)

        lr_scheduler.step(val_loss)
        if current_lr != get_lr(opt):
            if(verbose):
                print("Loading best model weights!")
            model.load_state_dict(best_model_wts)

        if(verbose):
            print(f"train loss: {train_loss:.6f}, dev loss: {val_loss:.6f}, accuracy: {100*val_metric:.2f}")
            print("-"*10)

    model.load_state_dict(best_model_wts)
    return model, loss_history, metric_history

解决方案 #

步骤1：修改自定义模型，记录最后卷积层输出 #

Grad-CAM需要获取模型最后一层卷积的特征图，因此需要在模型中新增变量保存该层输出：

class CNN_TUMOR(nn.Module):
    def __init__(self, params):
        super(CNN_TUMOR, self).__init__()
        # 保留原初始化代码...
        self.last_conv_output = None  # 新增：保存最后卷积层输出

    def forward(self,X):
        # 前三层卷积+池化逻辑不变
        X = F.relu(self.conv1(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv2(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv3(X))
        X = F.max_pool2d(X, 2, 2)
        
        # 记录最后卷积层ReLU后的输出（池化前）
        X = F.relu(self.conv4(X))
        self.last_conv_output = X  # 保存特征图
        X = F.max_pool2d(X, 2, 2)
        
        # 全连接层逻辑不变
        X = X.view(-1, self.num_flatten)
        X = F.relu(self.fc1(X))
        X = F.dropout(X, self.dropout_rate)
        X = self.fc2(X)
        return F.log_softmax(X, dim=1)

步骤2：实现Grad-CAM核心计算函数 #

该函数负责计算目标类的热力图，核心逻辑是通过梯度加权最后卷积层特征图：

def compute_grad_cam(model, input_tensor, target_class=None):
    model.eval()
    output = model(input_tensor)
    
    # 若未指定目标类，使用预测概率最高的类别
    if target_class is None:
        target_class = output.argmax(dim=1).item()
    
    # 计算目标类得分对最后卷积层的梯度
    model.zero_grad()
    target_score = output[0, target_class]
    target_score.backward()
    
    # 获取梯度并计算通道权重（空间维度平均）
    gradients = model.last_conv_output.grad[0]
    weights = torch.mean(gradients, dim=(1, 2))
    
    # 获取最后卷积层特征图
    features = model.last_conv_output[0]
    
    # 生成热力图：权重加权特征图求和后ReLU
    cam = torch.zeros(features.shape[1:], dtype=torch.float32).to(input_tensor.device)
    for i, w in enumerate(weights):
        cam += w * features[i]
    cam = F.relu(cam)
    
    # 归一化到0-1范围
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    return cam.cpu().numpy(), target_class

步骤3：实现可视化函数 #

将热力图与原图叠加，直观展示模型关注区域：

import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

def visualize_grad_cam(original_image, cam, target_class, class_names=["No Tumor", "Tumor"]):
    # 反归一化原图（根据你的数据预处理参数调整mean和std）
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img = original_image.permute(1,2,0).cpu().numpy()
    img = std * img + mean
    img = np.clip(img, 0, 1)
    
    # 将热力图resize到原图尺寸
    cam = np.array(Image.fromarray(cam).resize((img.shape[1], img.shape[0])))
    
    # 绘制三图对比：原图、热力图、叠加图
    fig, axs = plt.subplots(1, 3, figsize=(15,5))
    
    axs[0].imshow(img)
    axs[0].set_title(f"Original Image")
    axs[0].axis('off')
    
    axs[1].imshow(cam, cmap='jet')
    axs[1].set_title(f"Grad-CAM Heatmap\nClass: {class_names[target_class]}")
    axs[1].axis('off')
    
    # 叠加热力图与原图
    heatmap = np.uint8(255 * cam)
    heatmap = plt.cm.jet(heatmap)[:, :, :3]
    superimposed_img = heatmap * 0.4 + img * 0.6
    axs[2].imshow(superimposed_img)
    axs[2].set_title(f"Superimposed Image")
    axs[2].axis('off')
    
    plt.tight_layout()
    plt.show()

步骤4：完整使用示例 #

加载训练好的模型，对单张测试图片生成Grad-CAM：

# 初始化模型并加载权重
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cnn_model = CNN_TUMOR(params_model).to(device)
# 替换为你的权重文件路径
cnn_model.load_state_dict(torch.load("path/to/your/best_model_weights.pth"))

# 示例：生成随机测试张量（实际替换为你的预处理后测试图片）
input_img = torch.randn(1, 3, 256, 256).to(device)

# 计算Grad-CAM热力图
cam, target_class = compute_grad_cam(cnn_model, input_img)

# 可视化结果
visualize_grad_cam(input_img[0], cam, target_class)

内容的提问来源于stack exchange，提问作者AmirHossein