yolov5 Grad-CAM

使用Grad-CAM可视化YOLOv5模型的特征图
为了在YOLOv5中实现Grad-CAM进行可视化，可以利用yolov5-gradcam这一工具[^1]。此工具专门设计用于基于YOLOv5的Grad-CAM可视化。
修改models/yolo.py
具体来说，在Detect类中的forward函数内需做相应调整[^2]：
class Detect(nn.Module):
...

def forward(self, x):
# 原有代码...

# 新增代码：保存中间层输出以便后续计算梯度
self.feature_maps = []
for i in range(len(x)):
self.feature_maps.append(x[i].detach())

return output
登录后复制
上述改动允许捕获网络不同阶段产生的特征映射，这对于之后应用Grad-CAM至关重要。
添加新文件
还需创建几个新的Python脚本以支持Grad-CAM功能：

models/gradcam.py: 定义了处理Grad-CAM逻辑的核心组件。
import torch

class GradCAM(object):
"""Implementation of the gradient-weighted class activation mapping."""

def __init__(self, model, target_layer_names, use_cuda=True):
"""
Args:
model (torch.nn.Module): A PyTorch model.
target_layer_names (list[str]): List containing names of targeted layers.
use_cuda (bool): Whether to run on GPU or not.
"""
self.model = model.eval()
self.cuda = use_cuda and torch.cuda.is_available()

if self.cuda:
self.model.to('cuda')

# Register a hook function that will be called during backward pass
self.grads_val = {}
self.fmaps = {}
self.handles = []

for name, module in self.model.named_modules():
if any([name.endswith(tl_name) for tl_name in target_layer_names]):
handle_fmap = module.register_forward_hook(
lambda m, inp, outp: self._store_features(m, inp, outp))

handle_grad = module.register_backward_hook(
lambda m, grad_in, grad_out: self._store_gradients(m, grad_in, grad_out))

self.handles.extend((handle_fmap, handle_grad))

def _store_features(self, module, input_tensor, output_tensor):
self.fmaps[id(module)] = output_tensor.detach().cpu()

def _store_gradients(self, module, grad_input, grad_output):
self.grads_val[id(module)] = grad_output[0].detach().cpu()

def remove_hooks(self):
for hdl in self.handles:
hdl.remove()
登录后复制

models/yolov5_object_detector.py: 提供了一个接口来加载预训练好的YOLOv5权重并执行检测任务的同时也能够调用Grad-CAM分析特定目标类别对应的激活区域。

通过这些更改和新增加的功能模块，现在可以在运行预测的过程中同时获取到感兴趣的卷积层上的响应情况，并据此绘制出热力图表示哪些位置对于最终分类决策贡献较大。