【从零构建系列】YOLOv1

【从零构建系列】YOLOv1

（论文地址：https://arxiv.org/pdf/1506.02640.pdf）

模型介绍

以下为论文中的简述：

在YOLOv1提出之前，R-CNN系列算法在目标检测领域中独占鳌头。R-CNN系列检测精度高，但是由于其网络结构是双阶段（two-stage）的特点，使得它的检测速度不能满足实时性，饱受诟病。为了打破这一僵局，涉及一种速度更快的目标检测器是大势所趋。

2016年，Joseph Redmon、Santosh Divvala、Ross Girshick等人提出了一种单阶段（one-stage）的目标检测网络。它的检测速度非常快，每秒可以处理45帧图片，能够轻松地实时运行。由于其速度之快和其使用的特殊方法，作者将其取名为：You Only Look Once（也就是我们常说的YOLO的全称），并将该成果发表在了CVPR2016上，从而引起了广泛地关注。

YOLO的核心思想就是把目标检测转变成一个回归问题，利用整张图作为网络的输入，仅仅经过一个神经网络，得到bounding box（边界框）的位置及其所属的类别。

---

此处我们可以通过以下图点简单了解Faster-RCNN与YOLO：

总结：

1. YOLOv1没了RPN这个结构，更加专注于在backbone中得到最终的结果。（从理论上说应该是没有anchor结构定义，采用的是网格结构，在后来的YOLOv2结构中仍旧是anchor结构）
2. YOLOv1输出采用的是耦合头(分类+边框回归+置信度分数 在一起)

代码实现

要快速理解整个模型，必须先搞清楚它输入了什么，输出了什么（注意：目标检测中你不仅仅要知道我在上述图片中所表述的东西，也要知道正负样本划分，数据集的预处理），其次才是主干网络做了什么（这个时候你需要带入时代感，什么是时代感，就是说这个网络啥时候诞生的，这个时间点CNN、FCN、RCNN…都在干什么，为什么这样去思考，原因就是在于在后来的历史中主干网络其实有更多选择，比如编者在代码实现中主要使用MobileNetV3，此处就不解释了。最后你需要考虑如何训练，训练损失、优化器、是否启用衰减学习法等等…）

输入输出

输入数据展示（此处数据集来源于《动手学习深度学习》，目标为图中的香蕉,.jpg 格式）

注意：大小重要吗？ 不重要，反正都可以统一Resize，除非你输入的图像宽高比很离谱

常见的处理方式：1、Resize 2、Pad

输出数据展示（.txt）：

0 0.484375 0.152344 0.179688 0.164062

从左往右依次为“类别，真实框中心点坐标xy，宽，高”，都是按照占用原图宽高比例来计算的。（啊？你问为什么要这样做，因为如果是数值，那图片在进行数据增强的时候，岂不是惨了。）

然后如果图中有多个目标，则像如下内容生成即可。

0 0.484375 0.152344 0.179688 0.164062
0 0.484375 0.152344 0.179688 0.164062 # 此处仅仅是举例说明

数据差不多像这样

每一张图片对应它的标注

然后就是代码（此处不解释Dataset，DataLoader，看pytorch官网说明）：

import numpy as np
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import os
from PIL import Image
from torchvision.transforms import transforms


class YOLODataset(Dataset):
    def __init__(self, root_dir, s=7, b=2, c=1, transform=None):
        super().__init__()

        self.s = s  # 网格数量
        self.b = b  # bbox的数量
        self.c = c  # 分类数量
        self.transform = transform

        #  获取所有的图片地址
        images_path = os.path.join(root_dir, "images")
        image_names = os.listdir(images_path)
        self.image_urls = [os.path.join(images_path, name) for name in image_names]

        #  获取所有的图片标注信息
        labels_path = os.path.join(root_dir, "labels")
        label_names = os.listdir(labels_path)
        self.label_urls = [os.path.join(labels_path, name) for name in label_names]

    def __getitem__(self, index):
        image_path = self.image_urls[index]
        label_path = self.label_urls[index]

        with open(label_path, "r", encoding="utf-8") as f:
            bboxes = [[float(item) for item in line.strip().split()] for line in f.readlines()]
        bboxes = torch.tensor(bboxes)[:, [1, 2, 3, 4, 0]]

        img_pl = Image.open(image_path).convert("RGB")
        image = torch.from_numpy(np.array(img_pl)).permute([2, 0, 1])
        if self.transform:
            image = self.transform(image.float())

        target = torch.zeros((self.s, self.s, self.b * 5 + self.c))
        for bbox in bboxes:
            # 边框回归参数计算
            cx, cy, w, h, label = bbox.tolist()
            i = int(cx // (1 / self.s))
            j = int(cy // (1 / self.s))
            # 直接开方，一起前向传播。
            w_sqrt = w ** 0.5
            h_sqrt = h ** 0.5

            target[i, j, int(self.b * 5 + label)] = 1
            target[i, j, 0:5] = torch.tensor([cx, cy, w_sqrt, h_sqrt, 1.])

        return image, target

    def __len__(self):
        return len(self.image_urls)


def get_loader(root_dir, batch_size):
    transform = transforms.Compose([
        transforms.Resize((448, 448)),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    train_dataset = YOLODataset(os.path.join(root_dir, "train"), transform=transform)
    valid_dataset = YOLODataset(os.path.join(root_dir, "val"), transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=True)
    return train_loader, valid_loader


# ================= 测试数据集加载的代码 ==================
if __name__ == '__main__':
    train_loader, valid_loader = get_loader("../data", 5)
    for image, target in train_loader:
        print(image.shape)
        print(target.shape)
        break

模型结构

官方模型结构如下：

我并没有采用官方说的那个模型，而是直接使用了mobilenetv3作为主干网络。（原因我已说明）

此处的主干网络，论文已有说明是预训练过后的网络，因此采用了torchvision.models的权重。

from torch import nn
from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights


class YOLOv1(nn.Module):
    def __init__(self, s, b, c):
        super().__init__()

        self.s = s
        self.b = b
        self.c = c

        mobilenet = mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V1)
        self.backbone = mobilenet.features

        self.conv = nn.Sequential(
            nn.Conv2d(960, 1024, 1),
            nn.BatchNorm2d(1024),
            nn.ReLU(),
            nn.Conv2d(1024, 1024, 3, 2, 1),
            nn.BatchNorm2d(1024),
            nn.ReLU(),
        )

        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(1024 * 7 * 7, 1024),
            nn.ReLU(),
            nn.Linear(1024, int((self.b * 5 + self.c) * self.s * self.s))
        )

    def forward(self, x):
        x = self.backbone(x)
        x = self.conv(x)
        x = self.fc(x)
        return x.reshape(-1, int(self.b * 5 + self.c))

# ================= 测试模型结构的代码 ==================
if __name__ == '__main__':
    from torchsummary import summary
    model = YOLOv1(7, 2, 1)
    summary(model, (3, 448, 448), device="cpu")

模型训练

模型训练损失参照论文原图：

此处有几点内容你需要知道：

1、λ 参数是什么？λ coord 这是边框回归损失系数 设置为5， λ noobj 这是不包含对象框置信度损失系数 设置0.5

（为什么是5，0.5？ 因为正负样本不平衡，负样本数量远大于正样本数量）

2、正负样本又是如何划分？ 作者按照输出特征网格数量在原图中划分，如果预测中心在原图网格中则是正样本，反之则为负样本。如下图所示。

3、为什么w，h要开方，因为同样的均方差距离，对于小样本来说损失太大，大样本来说损失很小

from dataset.loader import get_loader
from backbone.mobilenet_yolo import YOLOv1
import torch
from torch import nn
from tqdm import tqdm

if __name__ == '__main__':
    train_loader, valid_loader = get_loader("./data", batch_size=5)
    s = 7
    b = 2
    c = 1
    device = torch.device("cuda")
    model = YOLOv1(s, b, c)
    model = model.to(device)

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

    lambda_coord = 5
    lambda_noobj = 0.5

    epochs = 1000
    for epoch in range(epochs):
        loss_list = []
        loop = tqdm(train_loader)
        for image, target in loop:
            image = image.to(device)
            target = target.reshape(-1, 5 * b + c)
            target = target.to(device)
            optimizer.zero_grad()
            predicts = model(image)

            # 包含目标对象的掩码
            mask_obj = target[:, 4] == 1
            # 不包含对象的掩码
            mask_noobj = target[:, 4] == 0

            # 位置损失
            loss1 = lambda_coord * criterion(predicts[mask_obj][:, 0:4], target[mask_obj][:, 0:4])
            # criterion(predicts[mask_obj][:, 5:9], target[mask_obj][:, 5:9])

            # 置信度损失
            loss2 = criterion(predicts[mask_obj][:, [4, 9]], target[mask_obj][:, [4, 9]]) + \
                    lambda_noobj * criterion(predicts[mask_noobj][:, [4, 9]], target[mask_noobj][:, [4, 9]])

            # 分类损失
            loss3 = criterion(predicts[mask_obj][:, 5 * b:], target[mask_obj][:, 5 * b:])

            loss = loss1 + loss2 + loss3
            loss.backward()
            optimizer.step()

            loop.set_postfix({"loss": f"{loss.item():.4f}"})
            loss_list.append(loss.item())

        avg_loss = sum(loss_list) / len(loss_list)
        print(f"epoch:{epoch + 1}/{epochs} -- loss:{avg_loss:.4f}")

性能表现

（1）优点

1. YOLO检测速度非常快。标准版本的YOLO可以每秒处理45张图像；YOLO的极速版本每秒可以处理150帧图像。这就意味着YOLO可以小于25毫秒延迟，实时地处理视频。对于欠实时系统，在准确率保证的情况下，YOLO速度快于其他方法。
2. YOLO实时检测的平均精度是其他实时监测系统的两倍。
3. 迁移能力强，能运用到其他新的领域（比如艺术品目标检测）。

（2）局限性

1. YOLO对相互靠近的物体，以及很小的群体检测效果不好，这是因为一个网格只预测了2个框，并且都只属于同一类。
2. 由于损失函数的问题，定位误差是影响检测效果的主要原因，尤其是大小物体的处理上，还有待加强。（因为对于小的bounding boxes，small error影响更大）。
   5张图像；YOLO的极速版本每秒可以处理150帧图像。这就意味着YOLO可以小于25毫秒延迟，实时地处理视频。对于欠实时系统，在准确率保证的情况下，YOLO速度快于其他方法。
3. YOLO实时检测的平均精度是其他实时监测系统的两倍。
4. 迁移能力强，能运用到其他新的领域（比如艺术品目标检测）。

（2）局限性

1. YOLO对相互靠近的物体，以及很小的群体检测效果不好，这是因为一个网格只预测了2个框，并且都只属于同一类。
2. 由于损失函数的问题，定位误差是影响检测效果的主要原因，尤其是大小物体的处理上，还有待加强。（因为对于小的bounding boxes，small error影响更大）。
3. YOLO对不常见的角度的目标泛化性性能偏弱。