【目标检测】YOLOv9理论解读与代码分析

前言

YOLO这个系列的故事已经很完备了，比如一些Decoupled-Head或者Anchor-Free等大的策略改动已经在YOLOv8固定下来，后面已经估计只有拿一些即插即用的tricks进行小改。

mmdetection框架的作者深度眸也在知乎上对“是否会有YOLOv9”这一观点发表看法：
在这里插入图片描述
然而，yolov9却还是在今年2月出来了，一作是中国台湾Academia Sinica的Chien-Yao Wang，和YOLOv4，v7是同一作者。

YOLOv9论文：https://arxiv.org/abs/2402.13616
YOLOv9仓库：https://github.com/WongKinYiu/yolov9

v9的改动和v8差别比较大，并不是在v8的基础上进行改进的，而是在作者之前的工作v7的基础上进行进一步改进。因此，要理解v9的相关理论，需要对v4和v7做一些回顾。

研究背景

YOLOv9的核心问题聚焦于多数方法都忽略了输入数据在前馈过程中可能具有不可忽略的信息损失量。

作者用下面一组图来说明该问题，下面是不同模型在深度空间中高维特征的可视化，以往的算法会造成信息丢失，导致可视化出现失真。
在这里插入图片描述
该篇论文，作者主要提出两点贡献：

1.可编程梯度信息（PGI）：通过辅助可逆分支生成可靠梯度，以便深层特征仍然可以保持足够的信息量。
2.基于ELAN设计了广义ELAN（GELAN）：在计算块选择上更加自由。

PGI：Programmable Gradient Information

可编程梯度信息：Programmable Gradient Information（PGI）这个概念乍一看有点拗口，看了不少解读论文，对此概念也只停留在论文的翻译，这里谈谈我的理解。

首先来回溯一下YOLOv7的这篇论文[1]，文中已经提及辅助训练的概念，如下图所示，图a是普通模型的一个正常检测输出流程，图b是在图a的基础上，在网络浅层特征直接引出Auxiliary head，损失只是根据浅层网络的特征来进行计算，这样有助于网络利用深层特征做损失的信息损失。
在这里插入图片描述
理解这一点后，再看YOLOv9里面的这张图：

这张图里放了四种架构：

图a是一个普通的PAN，和v7中的Normal model一样
图b是RevCol，其改动是在网络浅层旧加入一些本来只作用于深层的neck，用来储存浅层的特征信息，不过问题也很明显，内存开销太大(Heavy Cost)
图c是深度监督(Deep Supervision)，其思想是在网络浅层旧单独Copy一个检测头，这和v7中的图b思想一样。
图d就是yolov9提出的pgi思想，想法挺简单，一方面是继续保留Deep Supervision的设计，在浅层就搞一个检测头，另一个方面是单开一路，将原图单独塞入一个辅助可逆训练分支(Auxiliary Reversible Branch)，这其实类似于copy了一个主分支的backbone，蓝色的是原始的主分支，在主分支做neck部分的时候，一方面在浅层就直接做一个检测头，令一方面和原始一样，到深层再去检测。

因此，pgi并不是一种特定的网络结构，而是一种辅助训练的思想，它可以根据不同的网络特点进行结合，是自由的，并且参数可以随训练不断进行调整，按作者的话说就是可编程(这就是包装的艺术！）

由于是辅助训练，因此模型在训练阶段的参数量会变大很多，但对推理阶段并不会造成影响，推理仍使用主分支，因此对推理的速度影响不大。

GELAN： Generalized Efficient Layer Aggregation Network

GELAN这项工作并不是为了解决论文着眼的深度模型信息缺失的问题，而是对PGI的补充和优化。由于PGI的策略会导致网络过于庞大，计算成本过高，因此引入GELAN，用来缓解计算量。

对于GELAN的解读，还是需要重新追溯到YOLOv4，在YOLOv4中，作者采用了之前提到一种CSP(Cross Stage Partial)的架构思想，如下图所示。
在这里插入图片描述
图a是一个DenseNet架构，图b是CSPDenseNet的架构，从图中不难看出，CSP的思想就是将网络的特征图拆成两部分，一部分进入原始网络中做特征提取等操作，另一部分直接Concat到第一部分的输出。

YOLOv4就是将CSP思想应用到v3提出的Darknet中，变成CSPDarknet。下面是YOLOv5中对CSP部分的实现，其中不难看出，在实际使用中，通过两个卷积实现对特征图通道的对半拆分。

import torch
import torch.nn as nn


def autopad(k, p=None):  # kernel, padding
    # Pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


class Conv(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class BottleneckCSP(nn.Module):
    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(BottleneckCSP, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
        self.cv4 = Conv(2 * c_, c2, 1, 1)
        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
        self.act = nn.LeakyReLU(0.1, inplace=True)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

    def forward(self, x):
        y1 = self.cv3(self.m(self.cv1(x)))
        y2 = self.cv2(x)
        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
        
bcsp = BottleneckCSP(1,2)
print(bcsp)

在CSP之后，ELAN这项工作在此基础将该架构进一步发展。在YOLOv7[3]中，也曾对ELAN这个架构进行拓展，提出了一个拓展ELAN（E-ELAN），如下图所示：

在这里插入图片描述

如图c所示，ELAN的特点是跨层聚合，即同时聚合浅层特征和深层特征，主要是为了解决模型深度加深时梯度消失的问题，同时也提升特征的利用效率。

在这里插入图片描述
YOLOv9提出的GELAN并没有对ELAN做特别大的改动，只是将原本固定的一系列卷积层(conv)换成任意的Block(any block)，作者说这样可以降低模型复杂度，提升准确率和推理速度，个人觉得这一点其实挺牵强的…。

代码解读

YOLOv9这篇论文实际上可以拆分成两项工作，PGI和GELAN。因此，在代码方面，作者也进行了一系列拆分，整体框架仍然是采用v5那套，因此训练的方式和数据组织结构和v5通用。

这个仓库里面包含了两套模型，gelan和yolov9，yolov9等价于gelan+pgi，从作者给出的测试效果图来看，yolov9的数值明显要比gelan高出一点，因此，在使用yolov9代码时，完全可以忽略gelan。

在这里插入图片描述
作者在这里给了三个train和三个val，对应功能如下：

train.py ：训练GELAN模型
train_dual.py：训练带有一个辅助训练分支的GELAN模型
train_triple.py：训练带有2个辅助训练分支的GELAN模型
val、val_dual、val_triple三者功能对应上述三个train

根据YOLOv9的论文所述，YOLOv9模型是GELAN+1个辅助训练分支，因此训练和验证v9模型就使用train_dual.py、val_dual.py。

按照论文所述，YOLOv9共分四个版本，从小到大依次为小型(yolov9-s)、中型(yolov9-m)、紧凑型(yolov9-c)、扩展型(yolov9-e)，截至目前，该仓库只开源了后两者型号。

至于train_triple，估计是作者实验性的代码，论文里也未提到两个辅助检测分支的实验效果，效果估计不会有太大的提升。

另外，仓库里还有一些实验性的文件和yolov9无关，是作者令一项最新工作：YOLOR-Based Multi-Task Learning，这篇工作是想通过多个不同的任务，比如目标检测、实例分割、语义分割和图像描述来相互促进，这部分内容对应panoptic，不过目前这部分内容并不成熟，作者也没给出对应的数据集组织形式，使用时略过即可。

下面看一些除网络结构外的代码细节，比如，yolov9在辅助训练部分，加了一组检测头，相当于共有6个检测头，此代码对应DualDDetect，

class DualDDetect(nn.Module):
    # YOLO Detect head for detection models
    dynamic = False  # force grid reconstruction
    export = False  # export mode
    shape = None
    anchors = torch.empty(0)  # init
    strides = torch.empty(0)  # init

    def __init__(self, nc=80, ch=(), inplace=True):  # detection layer
        super().__init__()
        self.nc = nc  # number of classes
        self.nl = len(ch) // 2  # number of detection layers
        self.reg_max = 16
        self.no = nc + self.reg_max * 4  # number of outputs per anchor
        self.inplace = inplace  # use inplace ops (e.g. slice assignment)
        self.stride = torch.zeros(self.nl)  # strides computed during build

        c2, c3 = make_divisible(max((ch[0] // 4, self.reg_max * 4, 16)), 4), max((ch[0], min((self.nc * 2, 128))))  # channels
        c4, c5 = make_divisible(max((ch[self.nl] // 4, self.reg_max * 4, 16)), 4), max((ch[self.nl], min((self.nc * 2, 128))))  # channels
        self.cv2 = nn.ModuleList(
            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3, g=4), nn.Conv2d(c2, 4 * self.reg_max, 1, groups=4)) for x in ch[:self.nl])
        self.cv3 = nn.ModuleList(
            nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch[:self.nl])
        self.cv4 = nn.ModuleList(
            nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3, g=4), nn.Conv2d(c4, 4 * self.reg_max, 1, groups=4)) for x in ch[self.nl:])
        self.cv5 = nn.ModuleList(
            nn.Sequential(Conv(x, c5, 3), Conv(c5, c5, 3), nn.Conv2d(c5, self.nc, 1)) for x in ch[self.nl:])
        self.dfl = DFL(self.reg_max)
        self.dfl2 = DFL(self.reg_max)

    def forward(self, x):
        shape = x[0].shape  # BCHW
        d1 = []
        d2 = []
        for i in range(self.nl):
            d1.append(torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1))
            d2.append(torch.cat((self.cv4[i](x[self.nl+i]), self.cv5[i](x[self.nl+i])), 1))
        if self.training:
            return [d1, d2]
        elif self.dynamic or self.shape != shape:
            self.anchors, self.strides = (d1.transpose(0, 1) for d1 in make_anchors(d1, self.stride, 0.5))
            self.shape = shape

        box, cls = torch.cat([di.view(shape[0], self.no, -1) for di in d1], 2).split((self.reg_max * 4, self.nc), 1)
        dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
        box2, cls2 = torch.cat([di.view(shape[0], self.no, -1) for di in d2], 2).split((self.reg_max * 4, self.nc), 1)
        dbox2 = dist2bbox(self.dfl2(box2), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
        y = [torch.cat((dbox, cls.sigmoid()), 1), torch.cat((dbox2, cls2.sigmoid()), 1)]
        return y if self.export else (y, [d1, d2])

相比于修改之前的Detect，该新类最终输出两个检测头的结果，即d1和d2，之后做损失时，对应做两个检测头的损失：

class ComputeLoss:
    # Compute losses
    def __init__(self, model, use_dfl=True):
        device = next(model.parameters()).device  # get model device
        h = model.hyp  # hyperparameters

        # Define criteria
        BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device), reduction='none')

        # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
        self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0))  # positive, negative BCE targets

        # Focal loss
        g = h["fl_gamma"]  # focal loss gamma
        if g > 0:
            BCEcls = FocalLoss(BCEcls, g)

        m = de_parallel(model).model[-1]  # Detect() module
        self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02])  # P3-P7
        self.BCEcls = BCEcls
        self.hyp = h
        self.stride = m.stride  # model strides
        self.nc = m.nc  # number of classes
        self.nl = m.nl  # number of layers
        self.no = m.no
        self.reg_max = m.reg_max
        self.device = device

        self.assigner = TaskAlignedAssigner(topk=int(os.getenv('YOLOM', 10)),
                                            num_classes=self.nc,
                                            alpha=float(os.getenv('YOLOA', 0.5)),
                                            beta=float(os.getenv('YOLOB', 6.0)))
        self.assigner2 = TaskAlignedAssigner(topk=int(os.getenv('YOLOM', 10)),
                                            num_classes=self.nc,
                                            alpha=float(os.getenv('YOLOA', 0.5)),
                                            beta=float(os.getenv('YOLOB', 6.0)))
        self.bbox_loss = BboxLoss(m.reg_max - 1, use_dfl=use_dfl).to(device)
        self.bbox_loss2 = BboxLoss(m.reg_max - 1, use_dfl=use_dfl).to(device)
        self.proj = torch.arange(m.reg_max).float().to(device)  # / 120.0
        self.use_dfl = use_dfl

    def preprocess(self, targets, batch_size, scale_tensor):
        if targets.shape[0] == 0:
            out = torch.zeros(batch_size, 0, 5, device=self.device)
        else:
            i = targets[:, 0]  # image index
            _, counts = i.unique(return_counts=True)
            out = torch.zeros(batch_size, counts.max(), 5, device=self.device)
            for j in range(batch_size):
                matches = i == j
                n = matches.sum()
                if n:
                    out[j, :n] = targets[matches, 1:]
            out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor))
        return out

    def bbox_decode(self, anchor_points, pred_dist):
        if self.use_dfl:
            b, a, c = pred_dist.shape  # batch, anchors, channels
            pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.type(pred_dist.dtype))
            # pred_dist = pred_dist.view(b, a, c // 4, 4).transpose(2,3).softmax(3).matmul(self.proj.type(pred_dist.dtype))
            # pred_dist = (pred_dist.view(b, a, c // 4, 4).softmax(2) * self.proj.type(pred_dist.dtype).view(1, 1, -1, 1)).sum(2)
        return dist2bbox(pred_dist, anchor_points, xywh=False)

    def __call__(self, p, targets, img=None, epoch=0):
        loss = torch.zeros(3, device=self.device)  # box, cls, dfl
        feats = p[1][0] if isinstance(p, tuple) else p[0]
        feats2 = p[1][1] if isinstance(p, tuple) else p[1]
        
        pred_distri, pred_scores = torch.cat([xi.view(feats[0].shape[0], self.no, -1) for xi in feats], 2).split(
            (self.reg_max * 4, self.nc), 1)
        pred_scores = pred_scores.permute(0, 2, 1).contiguous()
        pred_distri = pred_distri.permute(0, 2, 1).contiguous()
        
        pred_distri2, pred_scores2 = torch.cat([xi.view(feats2[0].shape[0], self.no, -1) for xi in feats2], 2).split(
            (self.reg_max * 4, self.nc), 1)
        pred_scores2 = pred_scores2.permute(0, 2, 1).contiguous()
        pred_distri2 = pred_distri2.permute(0, 2, 1).contiguous()

        dtype = pred_scores.dtype
        batch_size, grid_size = pred_scores.shape[:2]
        imgsz = torch.tensor(feats[0].shape[2:], device=self.device, dtype=dtype) * self.stride[0]  # image size (h,w)
        anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)

        # targets
        targets = self.preprocess(targets, batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
        gt_labels, gt_bboxes = targets.split((1, 4), 2)  # cls, xyxy
        mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)

        # pboxes
        pred_bboxes = self.bbox_decode(anchor_points, pred_distri)  # xyxy, (b, h*w, 4)
        pred_bboxes2 = self.bbox_decode(anchor_points, pred_distri2)  # xyxy, (b, h*w, 4)

        target_labels, target_bboxes, target_scores, fg_mask = self.assigner(
            pred_scores.detach().sigmoid(),
            (pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
            anchor_points * stride_tensor,
            gt_labels,
            gt_bboxes,
            mask_gt)
        target_labels2, target_bboxes2, target_scores2, fg_mask2 = self.assigner2(
            pred_scores2.detach().sigmoid(),
            (pred_bboxes2.detach() * stride_tensor).type(gt_bboxes.dtype),
            anchor_points * stride_tensor,
            gt_labels,
            gt_bboxes,
            mask_gt)

        target_bboxes /= stride_tensor
        target_scores_sum = max(target_scores.sum(), 1)
        target_bboxes2 /= stride_tensor
        target_scores_sum2 = max(target_scores2.sum(), 1)

        # cls loss
        # loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum  # VFL way
        loss[1] = self.BCEcls(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
        loss[1] *= 0.25
        loss[1] += self.BCEcls(pred_scores2, target_scores2.to(dtype)).sum() / target_scores_sum2 # BCE

        # bbox loss
        if fg_mask.sum():
            loss[0], loss[2], iou = self.bbox_loss(pred_distri,
                                                   pred_bboxes,
                                                   anchor_points,
                                                   target_bboxes,
                                                   target_scores,
                                                   target_scores_sum,
                                                   fg_mask)
            loss[0] *= 0.25
            loss[2] *= 0.25
        if fg_mask2.sum():
            loss0_, loss2_, iou2 = self.bbox_loss2(pred_distri2,
                                                   pred_bboxes2,
                                                   anchor_points,
                                                   target_bboxes2,
                                                   target_scores2,
                                                   target_scores_sum2,
                                                   fg_mask2)
            loss[0] += loss0_
            loss[2] += loss2_

        loss[0] *= 7.5  # box gain
        loss[1] *= 0.5  # cls gain
        loss[2] *= 1.5  # dfl gain
        return loss.sum() * batch_size, loss.detach()  # loss(box, cls, dfl)

从这段代码不难看出，这里损失是将三部分损失box、cls和dfl损失对应相加，并且对不同类别的损失进行赋权，回归损失会占最主要的部分。

目前我跑了一下yolov9-c这个网络，模型大小约为98M，比起YOLOv5还是比较大的，加入辅助推理分支之后，模型大小比较大可以理解，不过对于部署推理时，应该可以做一些优化，比如把辅助推理的相关权重分支裁剪掉，这部分目前尚未实现。

总结

YOLOv9进一步延续了YOLOv7的工作，对于v7参数量太大的问题做了一些缓解。不过该网络还是比较偏学术性，目前模型的转换部署等均不够成熟，该论文的切入点比较小众，有点另辟蹊径的感觉，作者讲故事的思路和能力确实值得借鉴学习。

参考文献

[1] WANG C Y, BOCHKOVSKIY A, LIAO H Y. YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors[J].
[2] WANG C Y, MARK LIAO H Y, WU Y H, et al. CSPNet: A New Backbone that can Enhance Learning Capability of CNN[C/OL]//2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), Seattle, WA, USA. 2020. http://dx.doi.org/10.1109/cvprw50498.2020.00203. DOI:10.1109/cvprw50498.2020.00203.
[3] WANG C Y, BOCHKOVSKIY A, LIAO H Y. YOLOv7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors[J].