当前位置：首页 > article >正文

目标检测-YOLOv8

article 2024/11/14 20:52:49

YOLOv8

YOLOv8 是 YOLO 系列的最新版本，它在 YOLOv7 的基础上进行了多项改进，主要侧重于进一步提升推理速度、检测精度以及模型的通用性。与之前版本相比，YOLOv8 引入了新的技术和优化策略，使其在多个方面更具优势。

相比 YOLOv7 的改进与优势

更加轻量化的网络架构
YOLOv8 进一步简化了网络结构，引入了新型的 EfficientRep 主干网络，在保证性能的前提下进一步减少了参数量和计算量。这使得 YOLOv8 在推理速度上比 YOLOv7 更加优秀，尤其在移动设备和嵌入式系统上表现更为出色。
改进的特征融合模块
YOLOv8 优化了 PANet 和 FPN 的特征融合方式，通过引入新的 Path Aggregation Network（Rep-PAN），增强了多尺度特征的融合能力。相比 YOLOv7 的 PANet 结构，Rep-PAN 通过更高效的路径聚合策略，在不增加复杂度的情况下提升了目标检测的鲁棒性。
新型的损失函数
YOLOv8 引入了 CIoU（Complete Intersection over Union）的改进版本，提升了模型对目标框位置的精度。与 YOLOv7 相比，YOLOv8 的损失函数对边界框的形状、大小、中心点位置以及目标框的比例都有了更好的处理方式。
更强的 Anchor-Free 检测头
YOLOv8 继续优化了 Anchor-Free 机制，使其在检测不同尺度的目标时具有更好的泛化能力。通过动态分配锚框和改进的标签匹配策略，YOLOv8 比 YOLOv7 具有更高的定位精度和检测效率。
优化的训练策略
YOLOv8 引入了新的训练策略，如 Mosaic Augmentation、MixUp Augmentation 和 Self-Adversarial Training，使得模型在数据增强方面更加多样化，提升了模型的泛化能力。这些增强策略使 YOLOv8 能够在少量数据的情况下，依然保持较高的检测性能。

核心代码展示

下面展示 YOLOv8 的部分核心代码，包括新的 EfficientRep 主干网络和 CIoU 损失函数的实现。

import torch
import torch.nn as nn

# 1. EfficientRep Backbone 主干网络
class EfficientRep(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks):
        super(EfficientRep, self).__init__()
        self.blocks = nn.ModuleList([RepBlock(in_channels, out_channels) for _ in range(num_blocks)])
        
    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

# RepBlock: 用于 EfficientRep 中的基本构建块
class RepBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(RepBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        return self.relu(out)

# 2. YOLOv8 检测头
class YOLOv8Head(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(YOLOv8Head, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels // 2, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(in_channels // 2)
        self.conv2 = nn.Conv2d(in_channels // 2, num_classes, 1, bias=False)
    
    def forward(self, x):
        x = torch.relu(self.bn1(self.conv1(x)))
        return self.conv2(x)

# 3. CIoU 损失函数
class CIoULoss(nn.Module):
    def __init__(self):
        super(CIoULoss, self).__init__()

    def forward(self, pred_boxes, true_boxes):
        # 预测框和真实框的 (x1, y1, x2, y2)
        pred_x1, pred_y1, pred_x2, pred_y2 = torch.split(pred_boxes, 1, dim=1)
        true_x1, true_y1, true_x2, true_y2 = torch.split(true_boxes, 1, dim=1)

        # 计算交集
        inter_x1 = torch.max(pred_x1, true_x1)
        inter_y1 = torch.max(pred_y1, true_y1)
        inter_x2 = torch.min(pred_x2, true_x2)
        inter_y2 = torch.min(pred_y2, true_y2)
        inter_area = torch.clamp(inter_x2 - inter_x1, min=0) * torch.clamp(inter_y2 - inter_y1, min=0)

        # 计算预测框和真实框的面积
        pred_area = (pred_x2 - pred_x1) * (pred_y2 - pred_y1)
        true_area = (true_x2 - true_x1) * (true_y2 - true_y1)
        
        # 计算 IoU
        union_area = pred_area + true_area - inter_area
        iou = inter_area / torch.clamp(union_area, min=1e-6)
        
        # 计算边界框中心点差距
        pred_center_x = (pred_x1 + pred_x2) / 2
        pred_center_y = (pred_y1 + pred_y2) / 2
        true_center_x = (true_x1 + true_x2) / 2
        true_center_y = (true_y1 + true_y2) / 2
        center_distance = (pred_center_x - true_center_x)**2 + (pred_center_y - true_center_y)**2
        
        # 计算包围框对角线距离
        enclose_x1 = torch.min(pred_x1, true_x1)
        enclose_y1 = torch.min(pred_y1, true_y1)
        enclose_x2 = torch.max(pred_x2, true_x2)
        enclose_y2 = torch.max(pred_y2, true_y2)
        enclose_diagonal = (enclose_x2 - enclose_x1)**2 + (enclose_y2 - enclose_y1)**2

        # CIoU
        ciou = iou - center_distance / torch.clamp(enclose_diagonal, min=1e-6)

        return 1 - ciou

# 4. YOLOv8 整体网络结构
class YOLOv8(nn.Module):
    def __init__(self, num_classes):
        super(YOLOv8, self).__init__()
        self.backbone = EfficientRep(3, 64, 4)  # 主干网络
        self.neck = RepPAN()  # 特征金字塔网络
        self.head = YOLOv8Head(256, num_classes)  # 检测头

    def forward(self, x):
        x = self.backbone(x)  # 提取特征
        x = self.neck(x)  # 多尺度特征融合
        return self.head(x)  # 输出预测

代码解析

EfficientRep 主干网络
YOLOv8 的主干网络基于 EfficientRep，通过多个轻量化的 RepBlock 提取图像的多尺度特征。每个 RepBlock 包含两个卷积层和 BN 层，具有较低的计算开销，同时保持较强的特征提取能力。
YOLOv8Head
YOLOv8 的检测头通过进一步简化的卷积层进行预测，采用 1x1 卷积来减少参数量，同时结合了新的激活函数来提升训练收敛速度。
CIoU 损失函数
CIoU 相较于传统的 IoU，进一步考虑了预测框和真实框的中心点距离及边界框的形状差异，提升了目标框的定位精度。