萤火跑模型 | 多卡并行实现 YOLOv5 高性能训练

Vachel November 30, 2022

YOLO 系列模型在整个深度学习目标检测领域有举足轻重的地位，近些年模型性不断发展，工程落地中的应用也十分广泛。

最近来自美国的 Ultralytics 公司发布了第五代 YOLO 模型: YOLOv5。相比上代模型，YOLOv5 以更轻量的参数，更极致的推理速度，一经发布受到了学术界和工业界相关人士的广泛关注。

幻方 AI 最近在萤火集群上对该项工作进行了体验和优化。YOLOv5 没有相应的论文说明，仅开源了代码。我们通过自研的 3FS、数据集仓库、hfreduce、算子等对模型训练进行优化，实现了近 50% 的训练提速。本文为大家简要描述。

项目地址：https://github.com/ultralytics/yolov5

模型仓库：https://github.com/HFAiLab/yolov5

模型简介

YOLOv5 是一种单阶段目标检测算法，下图展示了其整体结构框图：

对于一个目标检测算法而言，我们通常可以将其划分为4个通用的模块，具体包括：输入端、基准网络、Neck网络与Head输出端，对应于上图中的4个红色模块。YOLOv5 在 YOLOv4 的基础上添加了一些新的改进思路，使其速度与精度都得到了极大的性能提升。主要的改进思路如下所示：

输入端：表示输入的图片，该阶段通常包含一个图像预处理阶段，即将输入图像缩放到网络的输入大小，并进行归一化等操作。YOLOv5 提出了一些改进思路，如在网络训练阶段，使用 Mosaic 数据增强操作提升模型的训练速度和网络的精度；提出一种自适应锚框计算与自适应图片缩放方法；
基准网络：通常是一些性能优异的分类器种的网络，该模块用来提取一些通用的特征表示。融合其它检测算法中的一些新思路。YOLOv5 中不仅使用了 CSPDarknet53 结构，而且使用了 Focus 结构作为基准网络；
Neck网络：目标检测网络在 BackBone 与最后的 Head 输出层之间往往会插入一些层，利用它可以进一步提升特征的多样性及鲁棒性。虽然 YOLOv5 同样用到了 SPP 模块、FPN+PAN 模块，但是实现的细节有些不同；
Head输出层：用来完成目标检测结果的输出。针对不同的检测算法，输出端的分支个数不尽相同，通常包含一个分类分支和一个回归分支。YOLOv5 输出层的锚框机制与 YOLOv4 相同，主要改进的是训练时的损失函数 GIOU_Loss，以及预测框筛选的 DIOU_nms。

模型实践

幻方 AI 基于 Ultralytics 开源的 YOLOv5 代码，进行多级多卡的训练支持，开源至 HFAI YOLOv5，并采用幻方一系列优化工具进行提速升级，包括 hfai.datasets 高速数据集仓库，hfreduce 并行训练，hfai.nn 高性能算子等功能。

数据集

YOLOv5 开源代码里给出的训练案例，使用的训练数据只有 128 张图片构成的 COCO 子数据集。数据规模较小，且 Dataloader 读取训练样本时需要在存储系统里打开大量的小文件，整体训练吞吐会受到很大影响。

源代码里用 Yaml 配置训练数据集，通过如下方式构建 Dataloader:

from torch.utils.data import DataLoader, Dataset

class LoadImagesAndLabels(Dataset):
  def __init__(self, ...):

    f = []  # image files
    for p in path if isinstance(path, list) else [path]:
        p = Path(p)  # os-agnostic
        if p.is_dir():  # dir
            f += glob.glob(str(p / '**' / '*.*'), recursive=True)
        elif p.is_file():  # file
            with open(p) as t:
                t = t.read().strip().splitlines()
                parent = str(p.parent) + os.sep
                f += [x.replace('./', parent) if x.startswith('./') else x for x in t]  # local to global path
        else:
            raise FileNotFoundError(f'{prefix}{p} does not exist')
    self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)

    ...

可以看到，单个样本直接从本地存储系统中读取，如果使用 COCO 全量数据集，则需要打开的小文件数量对存储系统的压力不可小觑。

这里，我们采用自研的高速数据集仓库，hfai.datasets，优化训练数据的格式。将大量的小文件通过 FFRecord 整合转化为若干个集中的大文件，并且可以使用 FFRecord 中优化的随机读和批量读，进行数据样本的高速随机访问。同时，hfai 数据集仓库封装了很多常用的公开数据集，并且采用幻方自研的高速存储 3FS 进行保存，可以实现极大的性能优化。改造非常简单，如下例所示：

from hfai.datasets import COCODetection
from ffrecord import FileReader
from ffrecord.torch import Dataset, DataLoader

class LoadImagesAndLabels(Dataset):
  def __init__(self, ...):
    self.data_dir = COCODetection('train').data_dir
    self.reader = FileReader(self.data_dir / "train2017.ffr", check_data=True)

    ...

通过上面两行代码，就可以轻松加载 hfai 数据集仓库中的 COCODetection 数据集，大幅提升训练样本的读取效率。更多的数据集和使用，可以阅读 hfai.datasets 官方文档。

hfai 训练改造

萤火提供了统一的训练管理平台，可以将大量的训练任务依据优先级负载均衡到不同的GPU上执行，充分利用起计算算力。对于 YOLOv5 的训练，我们可以加大数据并行，极大加速训练速度。操作上只需要加入如下几行代码：

引入 hfai
```
import hfai
```

使用 hfreduce, hfai.nn 替换代码中的参数

import hfai.nn as hfnn
from hfai.nn.parallel import DistributedDataParallel

def main(local_rank, args):

  ...

  model = hfnn.to_hfai(model)
  model = DistributedDataParallel(model.cuda(), device_ids=[local_rank])

  ...

使用 hfai.multiprocessing.spawn 启动多卡并行训练，绑定 numa 避免网络拥塞以提高训练效率

if __name__ == "__main__":
  args = parse_args()
  ngpus = torch.cuda.device_count()
  hfai.multiprocessing.spawn(main, (args,), nprocs=ngpus, bind_numa=True)

训练每步都轮询一次萤火集群，看当前的计算资源是否足够继续训练，做好断点的保存

start_epoch, start_step, model, optimizer, ... = load_ckpt()

def train(...):
  for epoch in range(start_epoch, epochs):
    for step, batch in enumerate(dataloader):
      if step < start_step:
        continue
      
      ...

      if rank == 0 and local_rank ==0 and hfai.client.receive_suspend_command():
        save_ckpt(model, epoch, step+1, optimizer, ...)
        hfai.go_suspend()

通过如上4步简单的改造，我们就可以基本利用起萤火集群的特性，对 YOLOv5 的训练进行优化提升。更多详细的代码改造细节，可以阅读我们在模型仓库中开源的代码。

训练结果

我们申请 4 个萤火计算节点进行并行训练。通过 hfai 命令行工具提交任务，测试萤火集群对 YOLOv5 训练加速的性能。如下图所示：

可以看到，左右两图分别展示了萤火集群特性（hfai）相比原始训练方式（raw）和不使用算子，hfreduce (hfai w/o nn&hfreduce) 的加速效果。如左图所示，完整的萤火集群特性相比原始训练方式有将近 50% 的训练提速，这得益于优秀的训练数据集管理和优良的算子与通信能力；如右图所示，在使用同样的数据加载模式下，hfreduce 和 hfai 优化算子可以进一步提升训练速度，充分利用显卡计算资源，整体获得将近 20% 的训练提速。

下面给大家展示一些萤火集群训练出的 YOLOv5 模型在目标检测任务上的效果：

体验总结

YOLOv5 作为目标检测领域的旗舰模型，受到了广泛的关注。其不提供论文直接开源代码，并且用 Yaml 配置文件来指导深度学习模型的构建，在深度学习领域中也是独树一帜。我们借助幻方萤火集群，很轻松地实现了集群特性与这套新的开发模式的打通，获得明显的加速效果，证明了萤火集群的易用性和实力。

综合体验打分如下：

研究指数：★★★

该模型是目标检测领域下的旗舰模型，成熟且已在工业界成功地落地实践。
开源指数：★★★★★

数据和代码都已开源，代码逻辑清晰可读性高。
门槛指数：★★★

样例数据量小，模型规模不大。代码开发模式新颖。
通用指数：★★

模型只针对目标检测场景，其他场景待验证。
适配指数：★★★★★

依赖简单，PyTorch 框架构建的模型，只需要修改几行代码就能在萤火集群上执行。

幻方 AI 紧跟 AI 研究的前沿浪潮，致力于用领先算力助力AI落地与价值创造，欢迎各方数据研究者与开发者们一同共建。