- 代码展示
import argparse import cv2 import numpy as np import torch from torch.autograd import Function from torchvision import models import torch torch.set_printoptions(profile="full") import torchvision from torchvision import transforms,datasets import time import os import numpy as np import pandas as pd from cv2 import cv2 from skimage import io from torch.utils.data import DataLoader,Dataset from PIL import Image import torch.nn as nn import torch.nn.functional as F import torch.backends.cudnn as cudnn from torch.optim import lr_scheduler import datetime class FeatureExtractor(): """ Class for extracting activations and registering gradients from targetted intermediate layers """ 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. """ 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 def preprocess_image(img): means = [0.485, 0.456, 0.406] stds = [0.229, 0.224, 0.225] preprocessed_img = img.copy()[:, :, ::-1] for i in range(3): preprocessed_img[:, :, i] = preprocessed_img[:, :, i] - means[i] preprocessed_img[:, :, i] = preprocessed_img[:, :, i] / stds[i] preprocessed_img = np.ascontiguousarray(np.transpose(preprocessed_img, (2, 0, 1))) preprocessed_img = torch.from_numpy(preprocessed_img) preprocessed_img.unsqueeze_(0) input = preprocessed_img.requires_grad_(True) return input def show_cam_on_image(img, mask): heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET) heatmap = np.float32(heatmap) / 255 cam = heatmap + np.float32(img) cam = cam / np.max(cam) cv2.imwrite("cam.jpg", np.uint8(255 * cam)) class GradCam: 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 def get_args(): parser = argparse.ArgumentParser() parser.add_argument('--use-cuda', action='store_true', default=True, help='Use NVIDIA GPU acceleration') parser.add_argument('--image-path', type=str, default='./S001C001P001R001A001_0_0.jpg', help='Input image path') args = parser.parse_args() args.use_cuda = args.use_cuda and torch.cuda.is_available() if args.use_cuda: print("Using GPU for acceleration") else: print("Using CPU for computation") return args def deprocess_image(img): """ see https://github.com/jacobgil/keras-grad-cam/blob/master/grad-cam.py#L65 """ img = img - np.mean(img) img = img / (np.std(img) + 1e-5) img = img * 0.1 img = img + 0.5 img = np.clip(img, 0, 1) return np.uint8(img*255) class DiffNet(nn.Module): def __init__(self): super(DiffNet,self).__init__() self.base_model = torchvision.models.resnet50(pretrained=True) # self.base_model.aux_logits = False self.flatten = nn.Flatten() self.base_model = nn.Sequential(*list(self.base_model.children())[:-2]) self.avgpooling = nn.AdaptiveAvgPool2d((1,1)) self.fc1 = nn.Linear(2048,1024) self.fc2 = nn.Linear(1024,512) # self.output1 = nn.Linear(512,2) self.output2 = nn.Linear(512,60) self.dropout = nn.Dropout(p=0.5) def forward(self,x): output_feature = self.base_model(x) output_feature = self.avgpooling(output_feature) output_feature = self.flatten(output_feature) # output_feature_reverse = ReverseLayerF.apply(output_feature,0.9) # output_feature_reverse = output_feature # dropout1 = self.dropout(output_feature_reverse) # fc1 = F.relu(self.fc1(dropout1)) # dropout2 = self.dropout(fc1) # fc2 = F.relu(self.fc2(dropout2)) # output1 = self.output1(fc2) # output1 = nn.LogSoftmax(output1)#rota_class dropout3 = self.dropout(output_feature) fc3 = self.fc1(dropout3) dropout4 = self.dropout(fc3) fc4 = self.fc2(dropout4) output2 = self.output2(fc4)#sk_class # output2 = nn.LogSoftmax(output2) return output1,output2,output_feature if __name__ == '__main__': """ python grad_cam.py <path_to_image> 1. Loads an image with opencv. 2. Preprocesses it for VGG19 and converts to a pytorch variable. 3. Makes a forward pass to find the category index with the highest score, and computes intermediate activations. Makes the visualization. """ args = get_args() # Can work with any model, but it assumes that the model has a # feature method, and a classifier method, # as in the VGG models in torchvision. # model = models.resnet50(pretrained=True) model = DiffNet() print(model) grad_cam = GradCam(model=model, feature_module=model.base_model, target_layer_names=["4"], use_cuda=args.use_cuda) img = cv2.imread(args.image_path, 1) img = np.float32(cv2.resize(img, (224, 224))) / 255 input = preprocess_image(img) # If None, returns the map for the highest scoring category. # Otherwise, targets the requested index. target_index = None mask = grad_cam(input, target_index) show_cam_on_image(img, mask) gb_model = GuidedBackpropReLUModel(model=model, use_cuda=args.use_cuda) print(model._modules.items()) gb = gb_model(input, index=target_index) gb = gb.transpose((1, 2, 0)) cam_mask = cv2.merge([mask, mask, mask]) cam_gb = deprocess_image(cam_mask*gb) gb = deprocess_image(gb) cv2.imwrite('gb.jpg', gb) cv2.imwrite('cam_gb.jpg', cam_gb)
Note：我是将自己的网络训练后保存模型，再加载未预训练的Resnet50模型，再加载自己主干网络的参数最后输出提取的特征，这样不用改动太多代码
比较便捷的一点就是若你的模型是单流的，直接把你的模型结构写上去，然后加载模型参数，设置feature_module和target_layer_name即可，若是多流的，则需要根据自己的实际情况debug，代码不长，比较好读（我没具体改过，应该不难）。还有一个trick是，若只想看主干网络提取的特征热图，Note的方法可以参考一下。
- 结果预览
待识别分类图像：
判别网络概率热图：
对抗网络概率热图
