YOLOv5详解(理论+代码)

YOLOv5详解(理论+代码)

YOLOv5于2020年首次提出,直至今日YOLOv5仍然在不断进行升级迭代。本文使用代码上5.0版本,yolov5根据精度和时间要求提供了YOLOv5n(6.0版本推出的)、YOLOv5s,YOLOv5m、YOLOv5l和YOLOv5x等不同版本(6.0之后还有一些针对小目标增加检测头的-P系列),这几个模型的结构基本一样,不同的是depth_multiple模型深度和width_multiple模型宽度,通过这两个参数来控制模型整体的参数量。下面给出主流模型参数量、推理时间和在COCO数据集上的mAP。

在这里插入图片描述
本文的使用代码为yolov5 5.0版本。代码下载地址: https://github.com/ultralytics/yolov5/archive/refs/tags/v5.0.zip

YOLOv5 网络结构

在这里插入图片描述

如上图所示,YOLOv5 5.0版本的网络结构主要有BackBone(New CSP-Darknet53)、Neck(FPN-PAN)和Head(YOLOv3 Head)三部分组成。具体的细节结构入下图所示:

在这里插入图片描述

Focus结构

Focus 可以看作是一种特殊的卷积操作,用于将输入的特征图进行压缩和重排,以减少模型的计算量和内存占用,并提高模型的精度和速度。具体操作是对输入的特征图进行间隔采样,然后在通道方向上进行堆叠,就得到了二倍下采样的特征图,在通过一个1x1的卷积调整通道数。虽然 Focus 重排操作可能会引入一些额外的计算和存储成本,但是通过 Focus 操作得到的二倍下采样特征图能够最大程度地保留输入特征图的所有信息。在 YOLOv5 中的应用证明了 Focus 的有效性。但是一些嵌入式推理框架对Focus的支持并不完善,在YOLOv5 6.0版本以后就用Stride=2的Conv结构代替了Focus。Focus 采样方式如图所示:
在这里插入图片描述

代码实现:

class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):  # ch_in, ch_out, kernel, stride, padding, groups
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
        # self.contract = Contract(gain=2)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))

Conv结构

Conv模块组成比较简单,有卷机+BN+SiLU激活函数构成,如果要提升网络在边缘设备上的推理速度,并可以接受小范围的精度损失,可以将SiLU激活函数用ReLU进行替换,对模型的推理速度有较大的提升。

代码实现

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)))

SPP结构

早期的目标检测算法,例如RCNN等利用Selective Search算法选取出图像中的一些候选区域,在通过CNN提取这些这个候选区域的特征,最后通过SVM进行分类算法和使用全连接层对Bounding-box的位置进行修正,提高定位精度。这就要求CNN输出的特征为度保持不变,因此输入CNN的候选区域尺寸是固定的,在检测各种大小不同的候选区域,需要经过resize、crop,或者warp等一系列操作,这都在一定程度上导致图片信息的丢失和变形,限制了识别精确度。

在这里插入图片描述

SPP-Net对这些网络中存在的缺点进行了改进,基本思想是:输入整张图像,提取出整张图像的特征图,如何在特征图中提取若干个候选区域(region proposal),再使用spatial pyramid pooling layer提取各个region proposal的特征。

在这里插入图片描述

spatial pyramid pooling layer 会将feature map 分成4x4、2x2、1x1个块(Spatial Bins),然后对每块进行全局池化,得到21x256维特征向量。所以SPP层的输出,不随输入尺度的变化而变化。这种方式也被后面的FastRCNN和FasterRCNN采用,在一些资料中被称为ROI Pooling。

虽然对于YOLOv5这种没有region proposal目标检测算法,并不存在region proposal尺度影响定位的问题,但是SPP这种通过不同尺度池化核聚合目标特征的方式依旧值得借鉴,Ultralytics 团队受到启发提出一种适合YOLO的空间池化金字塔(SpatialPyramid Pooling, SPP)结构。该结构首先利用 Conv模块提取特征得到输出,再利用 13×13、9×9、5×5 和 1×1 四个不同大小的池化核对 Conv 模块输出的特征图进行处理,然后将不同池化和的输出在通道方向上拼接在一起,再输入一个 Conv模块块细化特征得到 SPP 模块的输出。SPP 结构可以在基本不增加网络参数和计算量的情况下扩大模型的感受野,整合不同尺度的特征,提高模型检测能力,使得模型能够更好地适应复杂场景下的目标检测任务。
在这里插入图片描述

在YOLOv5 6.0版本中使用了SPP的改进模块SPPF,SPPF将SPP的并行改为了串行, 速度上提升了,精度上没有做过消融实验,暂时不清楚。

代码实现

SPP

class SPP(nn.Module):
    # Spatial pyramid pooling layer used in YOLOv3-SPP
    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

    def forward(self, x):
        x = self.cv1(x)
        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
      
#等效实现
class SPP(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool1 = nn.MaxPool2d(kernel_size = 5, stride = 1, padding=2)
        self.maxpool2 = nn.MaxPool2d(9, 1, padding=4)
        self.maxpool3 = nn.MaxPool2d(13, 1, padding=6)

    def forward(self, x):
        o1 = self.maxpool1(x)
        o2 = self.maxpool2(x)
        o3 = self.maxpool3(x)
        return torch.cat([x, o1, o2, o3], dim=1)

SPPF

class SPPF(nn.Module):
    # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
    def __init__(self, c1, c2, k=5):  # equivalent to SPP(k=(5, 9, 13))
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * 4, c2, 1, 1)
        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)

    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning
            y1 = self.m(x)
            y2 = self.m(y1)
            return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
#等效实现
class SPPF(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool = nn.MaxPool2d(5, 1, padding=2)

    def forward(self, x):
        o1 = self.maxpool(x)
        o2 = self.maxpool(o1)
        o3 = self.maxpool(o2)
        return torch.cat([x, o1, o2, o3], dim=1)

Neck部分结构

在这里插入图片描述

Bochkovskiy 等受PANet的启发在 YOLOv4 的Neck中使用 FPN 连接 PAN 结构形成改良的特征金字塔。FPN 用于将深层的语义特征传递到浅层,从而增强多个尺度上的语义表达,而 PAN 用于将浅层的定位信息传递到深层,增强多个尺度上的定位能力。大量实验证明了该结构的有效性,这一结构也被之后的YOLOv5系列网络所沿用。

Head部分结构

YOLOv5的头部结构与v3和v4都是一致的,通过一个1x1的卷积将Neck输出的三个特征层通道数都调整到(5+目标类别数)* 每个特征层的anchor数,YOLOv5一个特征层设置3个anchor,如果是COCO数据,有80个类别,那么特征通道数就应该是(5+80)* 3 = 255。然后进行解码将目标位置换算到输入图像上。其中5分别对应的是:预测框的中心点横坐标、纵坐标、宽度、高度和置信度。这里的置信度表示预测框的可信度,取值范围为( 0 , 1 ) ,值越大说明该预测框中越有可能存在目标。

Head中的3个检测层分别对应Neck中得到的3种不同尺寸的特征图。特征图上的每个网格都预设了3个不同宽高比的anchor,可以在特征图的通道维度上保存所有基于anchor先验框的位置信息和分类信息,用来预测和回归目标。

下面给出具体的预测特征图目标框回归流程:

在这里插入图片描述

如同所示,其中:

(bx, by, bw, bh)表示预测框的中心点x, y坐标、宽度和高度

(cx, cy)表示预测框中心点所在网格的左上角坐标

(tx, ty)表示预测框的中心点相对于网格左上角坐标的偏移量

(tw, th)表示预测框的宽高相对于anchor宽高的缩放比例

(pw,ph)表示表示先验框anchor的宽高

具体计算公式如下:
b x = ( 2 ∗ σ ( t x ) − 0.5 ) + c x b_x=(2* \sigma(t_x)-0.5) + c_x bx=(2σ(tx)0.5)+cx

b y = ( 2 ∗ σ ( t y ) − 0.5 ) + c y b_y=(2* \sigma(t_y)-0.5) + c_y by=(2σ(ty)0.5)+cy

b w = p w ∗ ( 2 ∗ σ ( t w ) ) 2 b_w=p_w*(2* \sigma(t_w))^2 bw=pw(2σ(tw))2

b h = p h ∗ ( 2 ∗ σ ( t h ) ) 2 b_h=p_h*(2* \sigma(t_h))^2 bh=ph(2σ(th))2

为了将预测框的中心点约束到当前网格中,使用Sigmoid函数处理偏移量,使预测的偏移值保持在(0,1)范围内。这样一来,根据目标框回归计算公式,预测框中心点坐标的偏移量保持在(−0.5,1.5)范围内,如上图紫色区域所示。预测框的宽度和高度对于anchor的放缩范围为(0,4)。

检测头代码

class Detect(nn.Module):
    stride = None  # strides computed during build
    export = False  # onnx export

    def __init__(self, nc=80, anchors=(), ch=()):  # detection layer
        super(Detect, self).__init__()
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        self.na = len(anchors[0]) // 2  # number of anchors
        self.grid = [torch.zeros(1)] * self.nl  # init grid
        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
        self.register_buffer('anchors', a)  # shape(nl,na,2)
        self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch)  # output conv

    def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
        self.training |= self.export
        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.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

                y = x[i].sigmoid()
                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i]) * self.stride[i]  # xy
                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                z.append(y.view(bs, -1, self.no))

        return x if self.training else (torch.cat(z, 1), x)

    @staticmethod
    def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()

非极大值抑制(NMS)

当我们得到对目标的预测后,一个目标通常会产生很多冗余的预测框。Non-maximum suppression(NMS)其核心思想在于抑制非极大值的目标,去除冗余,从而搜索出局部极大值的目标,找到最优值。

在我们对目标产生预测框后,往往会产生大量冗余的边界框,因此我们需要去除位置准确率低的边界框,保留位置准确率高的边界框。NMS的主要步骤为:

  1. 对于每个种类的置信度按照从大到小的顺序排序,选出置信度最高的边框。
  2. 遍历其余所有剩下的边界框,计算这些边界框与置信度最高的边框的IOU值。如果某一边界框和置信度最高的边框IOU阈值大于我们所设定的IOU阈值,这意味着同一个物体被两个重复的边界框所预测,则去掉这这个边框。
  3. 从未处理的边框中再选择一个置信度最高的值,重复第二步的过程,直到选出的边框不再有与它超过IOU阈值的边框。

代码实现:

def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
                        labels=()):
    """Runs Non-Maximum Suppression (NMS) on inference results

    Returns:
         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
    """

    nc = prediction.shape[2] - 5  # number of classes
    xc = prediction[..., 4] > conf_thres  # candidates

    # Settings
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
    max_det = 300  # maximum number of detections per image
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS

    t = time.time()
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        x = x[xc[xi]]  # confidence

        # Cat apriori labels if autolabelling
        if labels and len(labels[xi]):
            l = labels[xi]
            v = torch.zeros((len(l), nc + 5), device=x.device)
            v[:, :4] = l[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
            x = torch.cat((x, v), 0)

        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        # Detections matrix nx6 (xyxy, conf, cls)
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]

        # Filter by class
        if classes is not None:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

        # Apply finite constraint
        # if not torch.isfinite(x).all():
        #     x = x[torch.isfinite(x).all(1)]

        # Check shape
        n = x.shape[0]  # number of boxes
        if not n:  # no boxes
            continue
        elif n > max_nms:  # excess boxes
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence

        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
            # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
            weights = iou * scores[None]  # box weights
            x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
            if redundant:
                i = i[iou.sum(1) > 1]  # require redundancy

        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
            break  # time limit exceeded

    return output

损失函数

YOLOv5 损失函数由三部分组成:

  1. 分类损失(Classification Loss)采用的是BCE loss,只计算正样本的分类损失。

  2. 置信度损失 (Objectness Loss)采用的依然是BCE loss,用网络预测的目标边界框与GT Box的CIoU作为标签,计算的是所有样本的损失。

  3. 边界框回归损失(Box Regression Loss)采用的是CIoU loss,只计算正样本的分类损失。

L o s s = λ 1 L c l s + λ 2 L o b j + λ 3 L b o x Loss=\lambda_1L_{cls} + \lambda_2L_{obj} + \lambda_3L_{box} Loss=λ1Lcls+λ2Lobj+λ3Lbox

式中, λ 1 \lambda_1 λ1 λ 2 \lambda_2 λ2 λ 3 \lambda_3 λ3 为平衡参数,平衡三部分权重,默认通过工程data目录下的hyp.scratch.yaml文件进行配置cls、obj和box进行配置,默认取0.5,1.0和0.05。

分类损失

YOLOv5默认使用二元交叉熵函数来计算分类损失。二元交叉熵函数的定义为
L = − y l o g p − ( 1 − y ) l o g ( 1 − p ) = { − l o g p , y = 1 − l o g ( 1 − p ) , y = 0 L=-ylogp-(1-y)log(1-p)=\begin{cases} -logp && , y=1\\ -log(1-p) && ,y=0 \end{cases} L=ylogp(1y)log(1p)={logplog(1p),y=1,y=0
式中,yy为输入样本对应的标签(正样本为1,负样本为0),p为模型预测该输入样本为正样本的概率。假设:
p t = { p , y = 1 1 − p , y = 0 p_t=\begin{cases} p && ,y=1\\ 1 - p && ,y=0 \end{cases} pt={p1p,y=1,y=0
交叉熵函数的定义可简化为:
L = − l o g p t L=-logp_t L=logpt
YOLOv5使用二元交叉熵损失函数计算类别概率和目标置信度得分的损失,各个标签不是互斥的。YOLOv5使用多个独立的逻辑(logistic)分类器替换softmax函数,以计算输入属于特定标签的可能性。在计算分类损失进行训练时,对每个标签使用二元交叉熵损失。这也避免使用softmax函数而降低了计算复杂度。

置信度损失

每个预测框的置信度表示这个预测框的可靠程度,值越大表示该预测框越可靠,也表示越接近真实框。对于置信度标签,YOLO之前的版本认为所有存在目标的网格(正样本)对应的标签值均为1,其余网格(负样本)对应的标签值为0。但是这样带来的问题是有些预测框可能只是在目标的周围,而并不能精准预测框的位置。

YOLOv5的做法是,根据网格对应的预测框与真实框的CIoU作为该预测框的置信度标签。与计算分类损失一样,YOLOv5默认使用二元交叉熵函数来计算置信度损失。

同时,对于目标损失,在不同的预测特征层也给予了不同权重。
L o b j = 4.0 ∗ L o b j s m a l l + 1.0 ∗ L o b j m e d i u m + 0.4 ∗ L o b j l a r g e L_obj=4.0*L^{small}_{obj}+1.0*L^{medium}_{obj}+0.4*L^{large}_{obj} Lobj=4.0Lobjsmall+1.0Lobjmedium+0.4Lobjlarge
在源码中,针对预测小目标的预测特征层(大特征图)采用的权重是4.0,针对预测中等目标的预测特征层采用的权重是1.0,针对预测大目标的预测特征层(小特征图)采用的权重是0.4,作者说这是针对COCO数据集设置的超参数。

边界框回归损失

YOLOv5默认使用CIoU来计算边界框损失。CIoU基于DIoU得到,其中DIoU将预测框和真实框之间的距离,重叠率以及尺度等因素都考虑了进去,使得目标框回归变得更加稳定。DIoU的损失函数为:
L D I o U = 1 − I o U + ρ 2 ( b , b g t ) c 2 L_{DIoU}=1-IoU+ \frac{ \rho^2(b,b^{gt})}{c^2} LDIoU=1IoU+c2ρ2(b,bgt)
式中, b b b b g t b^{gt} bgt分别表示预测框和真实框的中心点,ρ表示两个中心点之间的欧式距离,c表示预测框和真实框的最小闭包区域的对角线距离,gt是ground truth缩写。如下图所示:

在这里插入图片描述

CIoU是在DIoU的惩罚项基础上添加了一个影响因子αv,这个因子将预测框的宽高比和真实框的宽高比考虑进去,即CIoU的损失计算公式为:
L C I o U = 1 − I o U + ρ 2 ( b , b g t ) c 2 + α v L_{CIoU}=1-IoU+ \frac{ \rho^2(b,b^{gt})}{c^2} + \alpha v LCIoU=1IoU+c2ρ2(b,bgt)+αv
式中, α \alpha α v v v 可以通过下列公式进行计算:
α = v ( 1 − I o U ) + v \alpha=\frac{v}{(1-IoU) + v} α=(1IoU)+vv

v = 4 π 2 ( a r c t a n w g t h g t − a r c t a n w h ) 2 v=\frac{4}{\pi^2}(arctan\frac{w^{gt}}{h^{gt}}-arctan\frac{w}{h})^2 v=π24(arctanhgtwgtarctanhw)2

式中, w g t w^{gt} wgt表示真实框的宽度 h g t h^{gt} hgt表示真实框的高度 w w w h h h 则表示预测框的宽度和高度。

损失计算代码

class ComputeLoss:
    # Compute losses
    def __init__(self, model, autobalance=False):
        super(ComputeLoss, self).__init__()
        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))
        BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))

        # 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, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)

        det = model.module.model[-1] if is_parallel(model) else model.model[-1]  # Detect() module
        self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02])  # P3-P7
        self.ssi = list(det.stride).index(16) if autobalance else 0  # stride 16 index
        self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
        for k in 'na', 'nc', 'nl', 'anchors':
            setattr(self, k, getattr(det, k))

    def __call__(self, p, targets):  # predictions, targets, model
        device = targets.device
        lcls, lbox, lobj = torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
        tcls, tbox, indices, anchors = self.build_targets(p, targets)  # targets

        # Losses
        for i, pi in enumerate(p):  # layer index, layer predictions
            b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
            tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj

            n = b.shape[0]  # number of targets
            if n:
                ps = pi[b, a, gj, gi]  # prediction subset corresponding to targets

                # Regression
                pxy = ps[:, :2].sigmoid() * 2. - 0.5
                pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
                pbox = torch.cat((pxy, pwh), 1)  # predicted box
                iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True)  # iou(prediction, target)
                lbox += (1.0 - iou).mean()  # iou loss

                # Objectness
                tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype)  # iou ratio

                # Classification
                if self.nc > 1:  # cls loss (only if multiple classes)
                    t = torch.full_like(ps[:, 5:], self.cn, device=device)  # targets
                    t[range(n), tcls[i]] = self.cp
                    lcls += self.BCEcls(ps[:, 5:], t)  # BCE

                # Append targets to text file
                # with open('targets.txt', 'a') as file:
                #     [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]

            obji = self.BCEobj(pi[..., 4], tobj)
            lobj += obji * self.balance[i]  # obj loss
            if self.autobalance:
                self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()

        if self.autobalance:
            self.balance = [x / self.balance[self.ssi] for x in self.balance]
        lbox *= self.hyp['box']
        lobj *= self.hyp['obj']
        lcls *= self.hyp['cls']
        bs = tobj.shape[0]  # batch size

        loss = lbox + lobj + lcls
        return loss * bs, torch.cat((lbox, lobj, lcls, loss)).detach()
  • 53
    点赞
  • 62
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
YoloV5 5.0版本的head结构主要包括三个部分:SPP结构、PAN结构以及最终的输出层。下面我会对这三个部分进行详细的介绍。 1. SPP结构 SPP(Spatial Pyramid Pooling)结构是在CNN中广泛应用的一种结构,用于对不同大小的输入图像进行池化。在YoloV5中,SPP结构被用于对特征图进行池化,以捕获不同尺度的特征信息。具体来说,SPP结构包含三个不同大小的池化层,每个池化层的步长为1,池化核大小分别为5x5、9x9和13x13。池化后的结果被连接在一起,形成一个多尺度的特征图,并送入后面的PAN结构中。 2. PAN结构 PAN(Path Aggregation Network)结构是YoloV5中引入的一种多尺度特征融合结构。它主要分为两个部分:上采样和特征融合。上采样通过将低分辨率的特征图进行插值操作,得到与高分辨率特征图相同大小的特征图。特征融合则是将不同尺度的特征图进行融合,以提高检测精度。具体来说,YoloV5中的PAN结构包含了5个卷积层和4个上采样层,其中第一个卷积层用于将SPP结构输出的多尺度特征图进行卷积操作,后面4个卷积层则用于处理上采样后的特征图。特征融合主要是通过上采样后的特征图和SPP结构输出的多尺度特征图进行concat操作得到的。 3. 输出层 最终的输出层是一个包含多个卷积层和全连接层的网络结构。它主要用于将PAN结构输出的特征图进行卷积操作,并预测目标的类别和位置信息。具体来说,输出层包含了3个卷积层,其中最后一个卷积层的输出通道数为3x(5+C),其中C为目标类别数。最终的输出是一个形状为[batch_size, num_anchors, grid_size, grid_size, 3x(5+C)]的张量,其中num_anchors表示每个网格点的先验框数量。在输出张量中,前4个数表示目标框的位置信息,第5个数表示目标框的置信度,后面的C个数则表示目标框所属的类别。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值