【PyTorch 实战3：YOLOv5检测模型】10min揭秘 YOLOv5 检测网络架构、工作原理以及pytorch代码实现（附代码实现！）

YOLOv5简介

YOLOv5（You Only Look Once, Version 5）是一种先进的目标检测模型，是YOLO系列的最新版本，由Ultralytics公司开发。该模型利用深度学习技术，能够在图像或视频中实时准确地检测出多个对象的位置及其类别，是计算机视觉领域的重要里程碑之一。下面将详细介绍YOLOv5的架构、性能、应用和未来发展方向。

一、与之前版本的对比

相较于YOLOv4和其他先前版本，YOLOv5在多个方面进行了改进和优化。首先，YOLOv5提供了更高的检测准确性，这得益于其新的模型架构以及对数据集和训练过程的细致调优。其次，YOLOv5在处理速度上也有所提升，这意味着它可以更快地对图像或视频进行检测和识别，使其在实时应用中更具竞争力。此外，YOLOv5还引入了一些新的特性，如自动批处理大小调整和更高效的图像处理流程，进一步提升了模型的性能和灵活性。

二、YOLOv5的架构

YOLOv5的架构基于深度卷积神经网络（CNN），采用了一种称为骨干网络（Backbone）的模块化设计。骨干网络通常由多个卷积层和池化层组成，用于从原始图像中提取特征。在YOLOv5中，采用了一种称为CSPDarknet的改进的骨干网络，它结合了Cross-Stage Partial连接（CSP）和Darknet53的优点，具有更好的特征提取能力和更快的训练速度。

除了骨干网络外，YOLOv5还包含了一系列用于检测和识别对象的头部（Head）模块。这些头部模块负责将从骨干网络中提取的特征映射转换为对象的边界框及其类别概率。YOLOv5采用了一种简单而有效的头部设计，包括多个卷积层和线性激活函数，以实现高效的对象检测。

YOLOv5网络整体框架图（来自大佬的：博客）
在这里插入图片描述
SPPF模块

SPPF模块代码实现：

import time
import torch
import torch.nn as nn


class SPP(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool1 = nn.MaxPool2d(5, 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)


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)


if __name__ == '__main__':
    input_tensor = torch.rand(8, 32, 16, 16)
    spp = SPP()
    sppf = SPPF()
    output1 = spp(input_tensor)
    output2 = sppf(input_tensor)

    print(torch.equal(output1, output2))

    t_start = time.time()
    for _ in range(100):
        spp(input_tensor)
    print(f"spp time: {time.time() - t_start}")

    t_start = time.time()
    for _ in range(100):
        sppf(input_tensor)
    print(f"sppf time: {time.time() - t_start}")

三、损失函数
这部分大佬的YOLOv5网络详解介绍得超级好，建议可以读读这篇博客。

四、代码实现

安装PyTorch和其他所需的Python库。具体可以看YOLOv5的仓库要求，然后按照以下步骤进行操作：

克隆YOLOv5仓库：

git clone https://github.com/ultralytics/yolov5.git

创建一个新的Python文件，比如yolov5_custom.py。
编写以下代码，实现YOLOv5模型的加载和预测功能：

import torch
from pathlib import Path
from models.yolo import Model
from utils.general import non_max_suppression, scale_coords

# 定义YOLOv5模型类
class YOLOv5:
    def __init__(self, weights='yolov5s.pt', img_size=640):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = Model(Path('yolov5s.yaml')).to(self.device)
        self.model.load_state_dict(torch.load(weights, map_location=self.device)['model'])
        self.model.eval()
        self.img_size = img_size

    # 预测函数
    def predict(self, images):
        img_size = self.img_size
        img = images.copy()  # 使用副本以防止原始图像被修改
        img = torch.from_numpy(img).to(self.device)
        img = img.float() / 255.0  # 像素值归一化到[0,1]
        if img.ndimension() == 3:
            img = img.unsqueeze(0)  # 添加batch维度 (1, H, W, C) -> (1, 1, H, W, C)
        img = img.permute(0, 3, 1, 2)  # 调整维度顺序 (1, 1, H, W, C) -> (1, C, H, W)

        # 图像尺寸调整
        img, _ = self.model.preprocess(img, img_size, img_size)

        # 预测
        pred = self.model(img)[0]  # 只获取第一个尺度的预测结果
        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
        
        # 将预测结果转换为图像上的坐标
        pred[0][:, :4] = scale_coords(img.shape[2:], pred[0][:, :4], images.shape).round()

        return pred

# 使用示例
if __name__ == "__main__":
    # 创建YOLOv5对象
    yolov5 = YOLOv5()

    # 加载图像
    image_path = 'example.jpg'
    image = torch.imread(image_path)

    # 进行预测
    results = yolov5.predict(image)

    # 打印预测结果
    for det in results[0]:
        xmin, ymin, xmax, ymax, conf, cls = det
        print(f"Class: {int(cls)}, Confidence: {conf:.2f}, BBox: [{xmin:.2f}, {ymin:.2f}, {xmax:.2f}, {ymax:.2f}]")