使用vit模型使用Grad-CAM的代码

当使用ViT（Vision Transformer）模型时，您可以使用以下代码来实现Grad-CAM：

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

class ViTGradCAM:
    def __init__(self, model):
        self.model = model
        self.feature_maps = None
        self.gradient = None

    def save_feature_maps(self, module, input, output):
        self.feature_maps = output.detach()

    def save_gradient(self, grad):
        self.gradient = grad[0].detach()

    def register_hooks(self):
        target_layer = self.model.blocks[-1]  # 修改为您希望可视化的目标层

        target_layer.register_forward_hook(self.save_feature_maps)
        target_layer.register_backward_hook(self.save_gradient)

    def generate_heatmap(self, input_image, target_class=None):
        self.model.zero_grad()
        output = self.model(input_image)

        if target_class is None:
            target_class = torch.argmax(output)

        output[0, target_class].backward()

        weights = F.adaptive_avg_pool2d(self.gradient, 1)
        heatmap = torch.mul(self.feature_maps, weights).sum(dim=1, keepdim=True)
        heatmap = F.relu(heatmap)
        heatmap /= torch.max(heatmap)
＊＊＊