grad-cam用于3D分割网络的代码修改——以及特征层非常规输出的解决方法

首先，我们看下chatgpt写的Gradcam框架。

import torch
import torchvision.models as models
from pytorch_grad_cam import GradCAM# Load your 3D segmentation model
model = models.segmentation3d()# Define the target layer
target_layer = model.conv3# Initialize the Grad-CAM class
grad_cam = GradCAM(model, target_layer)# Load your input tensor
input_tensor = torch.randn(1, 1, 32, 32, 32)# Get the Grad-CAM result
result = grad_cam(input_tensor)# Save the Grad-CAM result
torch.save(result, "grad_cam.pt")

为了个性化设计输出的cam，我们需要修改源码。也就是GradCAM()函数。
先解读一下原版

class GradCAM(BaseCAM):def __init__(self, model, target_layers, use_cuda=False,reshape_transform=None):super(GradCAM,self).__init__(model,target_layers,use_cuda,reshape_transform)def get_cam_weights(self,input_tensor,target_layer,target_category,activations,grads):#这个函数放到后面说了return np.mean(grads, axis=(2, 3))

可以看到并没有太多函数方法，我们打开基类BaseCAM()

    def __init__(self,model: torch.nn.Module,target_layers: List[torch.nn.Module],use_cuda: bool = False,reshape_transform: Callable = None,compute_input_gradient: bool = False,uses_gradients: bool = True) -> None:self.model = model.eval()self.target_layers = target_layersself.cuda = use_cudaif self.cuda:self.model = model.cuda()self.reshape_transform = reshape_transformself.compute_input_gradient = compute_input_gradientself.uses_gradients = uses_gradientsself.activations_and_grads = ActivationsAndGradients(self.model, target_layers, reshape_transform)# 实例化了ActivationsAndGradients类

看一下ActivationsAndGradients

class ActivationsAndGradients:# 自动调用__call__()函数，获取正向传播的特征层A和反向传播的梯度A'def __init__(self, model, target_layers, reshape_transform): # 传入模型参数，申明特征层的存储空间（self.activations）# 和回传梯度的存储空间（self.gradients）self.model = modelself.gradients = []self.activations = []self.reshape_transform = reshape_transformself.handles = []# 注意，上文指明目标网络层是是用列表存储的（target_layers = [model.down4]）# 源码设计的可以得到多层cam图# 这里注册了一个前向传播的钩子函数“register_forward_hook()”，其作用是# 在不改变网络结构的情况下获取某一层的输出，也就是获取正向传播的特征层for target_layer in target_layers:self.handles.append(target_layer.register_forward_hook(self.save_activation))# hasattr(object,name)返回值:如果对象有该属性返回True,否则返回False# 其作用是判断当前环境中是否存在该函数（解决版本不匹配的问题）if hasattr(target_layer, 'register_full_backward_hook'):self.handles.append(target_layer.register_full_backward_hook(self.save_gradient))else:# 注册反向传播的钩子函数“register_backward_hook”，用于存储反向传播过程中梯度图self.handles.append(target_layer.register_backward_hook(self.save_gradient))# 官方API文档对于register_forward_hook()函数有着类似的用法，# self.activations中存储了正向传播过程中的特征层def save_activation(self, module, input, output):activation = outputif self.reshape_transform is not None:activation = self.reshape_transform(activation)self.activations.append(activation.cpu().detach())# 与上述类似，只不过save_gradient()存储梯度信息，值得注意的是self.gradients的存储顺序def save_gradient(self, model, grad_input, grad_output):grad = grad_output[0]if self.reshape_transform is not None:grad = self.reshape_transform(grad)self.gradients = [grad.cpu().detach()] + self.gradients # 反向传播的梯度A’放在最前，目的是与特征层顺序一致def __call__(self, x):# 自动调用，会self.model(x)开始正向传播，注意此时并没有反向传播的操作self.gradients = []self.activations = []return self.model(x)def release(self):for handle in self.handles:handle.remove()# handle要及时移除掉，不然会占用过多内存

可以看到，ActivationsAndGradients类主要的功能是通过钩子函数获取正向传播的特征层和反向传播的梯度图，分别应用了register_forward_hook(hook)和register_backward_hook(hook)方法。这两类钩子函数的作用是自动获取某些中间变量，因为pytorch会自动舍弃图计算中间结果。比如自变量x，中间变量y和结果z，我们在反向传播过程中输出y的梯度时会提示“None”，这就是pytorch自动舍弃的结果，我们可以通过注册钩子函数将这些中间结果获取。

register_forward_hook(hook)：调用方法是“网络层结构.register_forward_hook(hook)”在相应的网络层结构正向传播时，获取其特征层，并执行自己定义好的hook函数中（其中包含model、input和output—输出特征层三个参数），来存储特征层信息。

register_backward_hook(hook)：同样在指定的网络层结构执行完.backward()函数后调用钩子函数hook(model, grad_input, grad_output)。model是指定的网络层结构，grad_input是该层网络的所有输入的梯度（bias）、该层网络输入变量x的梯度（weight）和网络权重的梯度（x）；而grad_output是指该层网络输出的梯度。

BaseCAM函数其他方法。先看2d是怎么处理的。

@staticmethoddef get_loss(output, target):loss = output # 直接将预测值作为Loss回传，本文展示的是语义分割的结果return loss@staticmethoddef get_cam_weights(grads): # GAP全局平均池化，得到大小为[B,C,1,1]# 因为我们输入一张图，所以B=1，C为特征层的通道数return np.mean(grads, axis=(2,3), keepdims=True)@staticmethoddef get_target_width_height(input_tensor):# 获取原图的高和宽width, height = input_tensor.size(-1), input_tensor.size(-2) return width, heightdef get_cam_image(self, activations, grads):# 将梯度图进行全局平均池化，weights大小为[1, C, 1, 1]，在通道上具有不同权重分布weights = self.get_cam_weights(grads) #对梯度图进行全局平均池化weighted_activations = weights[:, :, None, None] * activations #和原特征层加权乘cam = weighted_activations.sum(axis=1) # 在C维度上求和，得到大小为(1,h,w)return cam@staticmethoddef scale_cam_img(cam, target_size=None):# 将cam缩放到与原始图像相同的大小，并将其值缩放到[0,1]之间result = []for img in cam: # 因为传入的目标层（target_layers）可能为复数，所以一层一层看img = img - np.min(img) #减去最小值img = img / (1e-7 + np.max(img))if target_size is not None:img = cv.resize(img, target_size) # 注意：cv2.resize(src, (width, height))，width在height前result.append(img)result = np.float32(result)return resultdef compute_cam_per_layer(self, input_tensor):activations_list = [a.cpu().data.numpy() for a in self.activations_and_grads.activations] grads_list = [a.cpu().data.numpy() for a in self.activations_and_grads.gradients]target_size = self.get_target_width_height(input_tensor)cam_per_target_layer = []for layer_activations, layer_grads in zip(activations_list, grads_list):# 一张一张特征图和梯度对应着处理cam = self.get_cam_image(layer_activations, layer_grads)cam[cam<0] = 0 #ReLUscaled = self.scale_cam_img(cam, target_size) # 将CAM图缩放到原图大小，然后与原图叠加，这考虑到特征图可能小于或大于原图情况cam_per_target_layer.append(scaled[:, None, :]) # 在None标注的位置加入一个维度，相当于scaled.unsqueeze(1)，此时scaled大小为# [1,1,H,W]return cam_per_target_layerdef aggregate_multi_layers(self, cam_per_layer):cam_per_layer = np.concatenate(cam_per_layer, axis=1) # 在Channel维度进行堆叠，并没有做相加的处理cam_per_layer = np.maximum(cam_per_layer, 0) # 当cam_per_layer任意位置值小于0，则置为0result = np.mean(cam_per_layer, axis=1) # 在channels维度求平均，压缩这个维度，该维度返回为1# 也就是说如果最开始输入的是多层网络结构时，经过该方法会将这些网络结构# 在Channels维度上压缩，使之最后成为一张图return self.scale_cam_img(result)def __call__(self, input_tensor, target): # __init__()后自动调用__call__()方法# 这里的target就是目标的gt（双边缘）if self.use_cuda:input_tensor = input_tensor.cuda()# 正向传播的输出结果，创建ActivationsAndGradients类后调用__call__()方法，执行self.model(x)# 注意这里的output未经softmax，所以如果网络结构中最后的ouput不能经历激活函数output = self.activations_and_grads(input_tensor)[0]_output = output.detach().cpu()_output=_output.squeeze(0).squeeze(0)self.model.zero_grad()loss = self.get_loss(output, target)loss.backward(torch.ones_like(target), retain_graph=True)# 将输出结果作为Loss回传，记录回传的梯度图，# 梯度最大的说明在该层特征在预测过程中起到的作用最大，# 预测的部分展示出来就是整个网络预测时的注意力cam_per_layer = self.compute_cam_per_layer(input_tensor) # 计算每一层指定的网络结构中的cam图return self.aggregate_multi_layers(cam_per_layer) # 将指定的层结构中所有得到的cam图堆叠并压缩为一张图def __del__(self):self.activations_and_grads.release()def __enter__(self):return selfdef __exit__(self, exc_type, exc_value, exc_tb):self.activations_and_grads.release()if isinstance(exc_value, IndexError):# Handle IndexError here...print(f"An exception occurred in CAM with block: {exc_type}. Message: {exc_value}")return True

那么3d应该怎么修改呢，如果输出是[B,C,Z,Y,X]的话，可以一次输出Z中的一张图片。
首先在get_cam_weights函数中改成return np.mean(grads, axis=(3,4), keepdims=True)
然后get_cam_image中改成weighted_activations = weights[:, :, 10, :, :] * activations[:, :, 10, :, :]

那假如我的target_layers输出并不是单一数组，而是list或tuple怎么办呢？
那就在ActivationsAndGradients类中save_gradient和save_activation里面的output修改成你想输出的

官方代码
本文2d部分参考了：https://zhuanlan.zhihu.com/p/546875698

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！