# -*- coding:utf-8 -*- import argparse import cv2 import numpy as np import torch from torch.autograd import Function from torchvision import models ''' 调用关系 输入示例 model.layer4 --> feature module target_layer_names=["2"] --- target_layer_names GradCam-> ModelOutputs -> FeatureExtractor ''' class FeatureExtractor(): """ Class for extracting activations and registering gradients from targetted intermediate layers model：中间某个模块 target_layers:模块中的layer 1. 保存模块中target_layer的gradient 存储在self.gradients中 2. 将输入flow through模块，返回这中间的output的list和最终的output """ def __init__(self, model, target_layers): self.model = model self.target_layers = target_layers self.gradients = [] def save_gradient(self, grad): self.gradients.append(grad) def __call__(self, x): outputs = [] self.gradients = [] for name, module in self.model._modules.items(): x = module(x) if name in self.target_layers: x.register_hook(self.save_gradient) outputs += [x] return outputs, x class ModelOutputs(): """ Class for making a forward pass, and getting: 1. The network output. 2. Activations from intermeddiate targetted layers. 3. Gradients from intermeddiate targetted layers. 输入： 1. 模型 2. 模块 3. target layer 针对模型中的每一个模块 1. 如果这个模块是target 模块，则保留这个区域的feature map=== target_activations。 0 是进入的，-1是输出的 2. 如果是gap 则前向传递并reshape 3. 如果是其他模块 前向传递 """ def __init__(self, model, feature_module, target_layers): self.model = model self.feature_module = feature_module self.feature_extractor = FeatureExtractor(self.feature_module, target_layers) def get_gradients(self): return self.feature_extractor.gradients def __call__(self, x): target_activations = [] for name, module in self.model._modules.items(): if module == self.feature_module: target_activations, x = self.feature_extractor(x) elif "avgpool" in name.lower(): x = module(x) x = x.view(x.size(0),-1) else: x = module(x) return target_activations, x class GradCam: ''' call: 1. 返回特定模块的features 以及整个网络的output 2. 根据output计算预测的label 3. 计算target类别的loss 4. 清空feature module部分的grad 5. 模型zero grad 6. loss反传 7. 计算gradient 8. 取出想要的feature map 9. 根据grad value计算每个channel的权值 10. 对feature map进行加权 11. resize等等 ''' def __init__(self, model, feature_module, target_layer_names, use_cuda): self.model = model self.feature_module = feature_module self.model.eval() self.cuda = use_cuda if self.cuda: self.model = model.cuda() self.extractor = ModelOutputs(self.model, self.feature_module, target_layer_names) def forward(self, input): return self.model(input) def __call__(self, input, index=None): if self.cuda: features, output = self.extractor(input.cuda()) else: features, output = self.extractor(input) if index == None: index = np.argmax(output.cpu().data.numpy()) one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32) one_hot[0][index] = 1 one_hot = torch.from_numpy(one_hot).requires_grad_(True) if self.cuda: one_hot = torch.sum(one_hot.cuda() * output) else: one_hot = torch.sum(one_hot * output) self.feature_module.zero_grad() self.model.zero_grad() one_hot.backward(retain_graph=True) grads_val = self.extractor.get_gradients()[-1].cpu().data.numpy() target = features[-1] target = target.cpu().data.numpy()[0, :] weights = np.mean(grads_val, axis=(2, 3))[0, :] cam = np.zeros(target.shape[1:], dtype=np.float32) for i, w in enumerate(weights): cam += w * target[i, :, :] cam = np.maximum(cam, 0) cam = cv2.resize(cam, input.shape[2:]) cam = cam - np.min(cam) cam = cam / np.max(cam) return cam class GuidedBackpropReLU(Function): @staticmethod def forward(self, input): positive_mask = (input > 0).type_as(input) output = torch.addcmul(torch.zeros(input.size()).type_as(input), input, positive_mask) self.save_for_backward(input, output) return output @staticmethod def backward(self, grad_output): input, output = self.saved_tensors grad_input = None positive_mask_1 = (input > 0).type_as(grad_output) positive_mask_2 = (grad_output > 0).type_as(grad_output) grad_input = torch.addcmul(torch.zeros(input.size()).type_as(input), torch.addcmul(torch.zeros(input.size()).type_as(input), grad_output, positive_mask_1), positive_mask_2) return grad_input class GuidedBackpropReLUModel: def __init__(self, model, use_cuda): self.model = model self.model.eval() self.cuda = use_cuda if self.cuda: self.model = model.cuda() def recursive_relu_apply(module_top): for idx, module in module_top._modules.items(): recursive_relu_apply(module) if module.__class__.__name__ == 'ReLU': module_top._modules[idx] = GuidedBackpropReLU.apply # replace ReLU with GuidedBackpropReLU recursive_relu_apply(self.model) def forward(self, input): return self.model(input) def __call__(self, input, index=None): if self.cuda: output = self.forward(input.cuda()) else: output = self.forward(input) if index == None: index = np.argmax(output.cpu().data.numpy()) one_hot = np.zeros((1, output.size()[-1]), dtype=np.float32) one_hot[0][index] = 1 one_hot = torch.from_numpy(one_hot).requires_grad_(True) if self.cuda: one_hot = torch.sum(one_hot.cuda() * output) else: one_hot = torch.sum(one_hot * output) # self.model.features.zero_grad() # self.model.classifier.zero_grad() one_hot.backward(retain_graph=True) output = input.grad.cpu().data.numpy() output = output[0, :, :, :] return output