Spaces:
Sleeping
Sleeping
| """ | |
| Grad-CAM implementation for visualization | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| import numpy as np | |
| class GradCAM: | |
| """Grad-CAM实现""" | |
| def __init__(self, model, target_layer): | |
| self.model = model | |
| self.target_layer = target_layer | |
| self.gradients = None | |
| self.activations = None | |
| self.forward_handle = None | |
| self.backward_handle = None | |
| # 强制开启梯度（确保梯度能回传） | |
| for param in self.model.parameters(): | |
| param.requires_grad = True | |
| # 注册hook | |
| self.hook_layers() | |
| def hook_layers(self): | |
| """注册forward和backward hook""" | |
| if self.forward_handle is not None: | |
| self.forward_handle.remove() | |
| if self.backward_handle is not None: | |
| self.backward_handle.remove() | |
| def backward_hook(module, grad_in, grad_out): | |
| """标准的backward hook签名：(module, grad_in, grad_out)""" | |
| if grad_out[0] is not None: | |
| self.gradients = grad_out[0] | |
| def forward_hook(module, input, output): | |
| """forward hook保存激活值""" | |
| self.activations = output | |
| self.forward_handle = self.target_layer.register_forward_hook(forward_hook) | |
| self.backward_handle = self.target_layer.register_backward_hook(backward_hook) | |
| def generate_cam(self, input_tensor, class_idx=None): | |
| """生成CAM热力图""" | |
| self.gradients = None | |
| self.activations = None | |
| # 设置为train模式以允许梯度回传 | |
| self.model.train() | |
| # 禁用dropout和batchnorm的随机性 | |
| for m in self.model.modules(): | |
| if isinstance(m, (torch.nn.BatchNorm2d, torch.nn.Dropout)): | |
| m.eval() | |
| input_tensor = input_tensor.requires_grad_(True) | |
| # Forward pass | |
| output = self.model(input_tensor) | |
| # 如果是SGD-Net，取分类输出 | |
| if isinstance(output, tuple): | |
| output = output[0] | |
| # 如果没有指定类别，使用预测类别 | |
| if class_idx is None: | |
| class_idx = output.argmax(dim=1).item() | |
| # Backward pass | |
| self.model.zero_grad() | |
| output[0, class_idx].backward(retain_graph=True) | |
| # 如果gradients仍然是None或全0，使用fallback方法 | |
| if self.gradients is None or (hasattr(self.gradients, 'abs') and self.gradients.abs().max().item() < 1e-8): | |
| if self.activations is not None: | |
| if len(self.activations.shape) == 4: | |
| activations = self.activations[0] # (C, H, W) | |
| else: | |
| activations = self.activations | |
| cam = torch.mean(activations.abs(), dim=0) # (H, W) | |
| cam_min = cam.min() | |
| cam_max = cam.max() | |
| if cam_max > cam_min + 1e-8: | |
| cam = (cam - cam_min) / (cam_max - cam_min) | |
| else: | |
| cam = torch.ones_like(cam) * 0.5 | |
| return cam.detach().cpu().numpy() | |
| if self.gradients is None or self.activations is None: | |
| return np.ones((8, 8)) * 0.5 | |
| # 计算权重（全局平均池化梯度） | |
| if len(self.gradients.shape) == 4: | |
| gradients = self.gradients[0] # (C, H, W) | |
| else: | |
| gradients = self.gradients | |
| if len(self.activations.shape) == 4: | |
| activations = self.activations[0] # (C, H, W) | |
| else: | |
| activations = self.activations | |
| # 计算权重（对每个通道的梯度进行全局平均池化） | |
| weights = torch.mean(gradients, dim=(1, 2), keepdim=True) # (C, 1, 1) | |
| # 加权求和 | |
| cam = torch.sum(weights * activations, dim=0) # (H, W) | |
| cam = F.relu(cam) # ReLU（只保留正梯度） | |
| # 移除左上角异常高值（针对SGD-Net） | |
| h, w = cam.shape | |
| corner_h = h // 3 | |
| corner_w = w // 3 | |
| corner_region = cam[:corner_h, :corner_w] | |
| rest_region = torch.cat([ | |
| cam[corner_h:, :].flatten(), | |
| cam[:, corner_w:].flatten() | |
| ]) | |
| if rest_region.numel() > 0: | |
| rest_median = torch.median(rest_region) | |
| rest_std = torch.std(rest_region) | |
| corner_max = torch.max(corner_region) | |
| corner_mean = torch.mean(corner_region) | |
| if corner_max > rest_median + 1.5 * rest_std or corner_mean > rest_median + 0.5 * rest_std: | |
| for i in range(corner_h): | |
| for j in range(corner_w): | |
| dist_to_edge = min(i, j) | |
| weight = max(0.0, 1.0 - dist_to_edge / max(corner_h, corner_w) * 2.0) | |
| weight = min(1.0, weight) | |
| cam[i, j] = cam[i, j] * (1 - weight) + rest_median * weight | |
| # 归一化 | |
| cam_min = cam.min() | |
| cam_max = cam.max() | |
| if cam_max > cam_min + 1e-8: | |
| cam = (cam - cam_min) / (cam_max - cam_min) | |
| else: | |
| cam = torch.ones_like(cam) * 0.5 | |
| cam = torch.clamp(cam, 0, 1) | |
| return cam.detach().cpu().numpy() | |
| def __del__(self): | |
| """清理hook""" | |
| if self.forward_handle is not None: | |
| self.forward_handle.remove() | |
| if self.backward_handle is not None: | |
| self.backward_handle.remove() | |
| def get_target_layer(model): | |
| """获取目标层用于Grad-CAM（针对ResNetSGDNet）""" | |
| backbone = model.backbone | |
| # ResNet - 取layer4的最后一个block的最后一个conv | |
| if hasattr(backbone, 'layer4'): | |
| layer4 = backbone.layer4 | |
| if hasattr(layer4, '__len__') and len(layer4) > 0: | |
| last_block = layer4[-1] | |
| if hasattr(last_block, 'conv3'): | |
| return last_block.conv3 | |
| elif hasattr(last_block, 'conv2'): | |
| return last_block.conv2 | |
| else: | |
| return last_block | |
| else: | |
| return layer4 | |
| else: | |
| # 默认取最后一个模块 | |
| modules = list(backbone.children()) | |
| return modules[-1] if modules else None | |