反卷积可视化：使用ResNet18预训练模型做推理

import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from torchvision.transforms.functional import to_pil_image
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import cv2

# 加载预训练的 ResNet18 模型
resnet18 = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
resnet18.eval()
print("-------GradCAM----------")

class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradient = None
        self.activation = None
        self.hook_handles = []

        # 注册前向传播和反向传播的钩子
        self._register_hooks()

    def _register_hooks(self):
        def forward_hook(module, input, output):
            self.activation = output

        def backward_hook(module, grad_input, grad_output):
            self.gradient = grad_output[0]

        # 注册钩子
        target_layer = self._get_target_layer()
        self.hook_handles.append(target_layer.register_forward_hook(forward_hook))
        self.hook_handles.append(target_layer.register_backward_hook(backward_hook))

    def _get_target_layer(self):
        # 获取目标层
        if isinstance(self.target_layer, str):
            for name, module in self.model.named_modules():
                if name == self.target_layer:
                    return module
        elif isinstance(self.target_layer, nn.Module):
            return self.target_layer
        else:
            raise ValueError("Invalid target layer type")

    def generate_cam(self, input_tensor, target_class):
        # 前向传播
        output = self.model(input_tensor)
        self.model.zero_grad()

        # 计算目标类别的梯度
        target = output[0, target_class]
        target.backward()

        # 计算 Grad-CAM
        gradient = self.gradient[0].cpu().data.numpy()
        activation = self.activation[0].cpu().data.numpy()

        weights = np.mean(gradient, axis=(1, 2))  # 计算权重
        cam = np.zeros(activation.shape[1:], dtype=np.float32)

        # 加权特征图
        for i, w in enumerate(weights):
            cam += w * activation[i]

        # ReLU 操作
        cam = np.maximum(cam, 0)
        cam = cv2.resize(cam, input_tensor.shape[2:])  # 缩放到输入图像大小
        cam = cam - np.min(cam)
        cam = cam / np.max(cam)
        return cam

    def remove_hooks(self):
        # 移除钩子
        for handle in self.hook_handles:
            handle.remove()