最新版的YOLOv5输出GradCAM热力图

文章目录

• 一、怎么输出GradCAM热力图？
• 二、使用步骤
  • 1.引入库
  • 2.读入数据
• 总结

一、YOLOV5输出热力图？

由于GitHub上的YoLo项目一直在更新，导致旧版本添加GradCAM热力图的方式不再使用，为了保证成功输出热力图，在此文进行详细介绍。

二、使用步骤

---修改部分
   ---model/yolo.py(Detect类中的forward函数)

---添加部分
   ---model/gradcam.py
   ---model/yolov5_object_detector.py
   ---main_gradcam.py
 

修改yolo代码如下（示例）：

修改model/yolo.py文件中的Detect类中的forward函数
添加如下四条语句：
logits_ = []
logits = x[i][..., 5:]
logits_.append(logits.view(bs, -1, self.no - 5))
out = (torch.cat(z, 1), torch.cat(logits_, 1), x)

 def forward(self, x):
        z = []  # inference output

        logits_ = []  # 修改---1
        for i in range(self.nl):
            x[i] = self.m[i](x[i])  # conv
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

                logits = x[i][..., 5:]  # 修改---2
                if isinstance(self, Segment):  # (boxes + masks)
                    xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
                    xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)
                else:  # Detect (boxes only)
                    xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
                    xy = (xy * 2 + self.grid[i]) * self.stride[i]  # xy
                    wh = (wh * 2) ** 2 * self.anchor_grid[i]  # wh
                    y = torch.cat((xy, wh, conf), 4)
                z.append(y.view(bs, self.na * nx * ny, self.no))  

                logits_.append(logits.view(bs, -1, self.no - 5))  # 修改---3
        # return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(logits_, 1), x)
        return x if self.training else (torch.cat(z, 1), torch.cat(logits_, 1), x)   #修改4

2.

在models目录下添加gradcam.py文件

代码如下：

import time
import torch
import torch.nn.functional as F


def find_yolo_layer(model, layer_name):
    """Find yolov5 layer to calculate GradCAM and GradCAM++

    Args:
        model: yolov5 model.
        layer_name (str): the name of layer with its hierarchical information.

    Return:
        target_layer: found layer
    """
    hierarchy = layer_name.split('_')
    target_layer = model.model._modules[hierarchy[0]]

    for h in hierarchy[1:]:
        target_layer = target_layer._modules[h]
    return target_layer


class YOLOV5GradCAM:
    # 初始化，得到target_layer层
    def __init__(self, model, layer_name, img_size=(640, 640)):
        self.model = model
        self.gradients = dict()
        self.activations = dict()

        def backward_hook(module, grad_input, grad_output):
            self.gradients['value'] = grad_output[0]
            return None

        def forward_hook(module, input, output):
            self.activations['value'] = output
            return None

        target_layer = find_yolo_layer(self.model, layer_name)
        # 获取forward过程中每层的输入和输出，用于对比hook是不是正确记录
        target_layer.register_forward_hook(forward_hook)
        target_layer.register_full_backward_hook(backward_hook)
        device = 'cuda' if next(self.model.model.parameters()).is_cuda else 'cpu'
        self.model(torch.zeros(1, 3, *img_size, device=device))

    def forward(self, input_img, class_idx=True):
        """
        Args:
            input_img: input image with shape of (1, 3, H, W)
        Return:
            mask: saliency map of the same spatial dim