YOLO-V5 系列算法和代码解析（六）—— 分布式训练

最新推荐文章于 2025-03-25 11:43:02 发布

X_Imagine

最新推荐文章于 2025-03-25 11:43:02 发布

阅读量3.2k

点赞数 1

分类专栏： # YOLO-V5代码解读文章标签：算法分布式深度学习 YOLO-V5 分布式训练 DP DDP

本文链接：https://blog.csdn.net/kxh123456/article/details/127571417

版权

YOLO-V5代码解读专栏收录该内容

10 篇文章

订阅专栏

文章目录

预备知识

为了更好的理解分布式相关的内容，需要提前熟悉一些基本概念和特定的名称。分布式训练涉及到设备端（CPU，GPU），算法端（梯度更新，反向传播）等基本概念，具体陈述如下：

基本概念

网络参数：【w，b】是训练过程中，需要更新的参数，也称为权重和偏置；w 的更新公式如下，
$w^*=w-\alpha\cdot \frac{\partial E}{\partial w} \tag1$
注： $w^*$ 是待更新的权重参数， $w$ 是上一次迭代的权重参数， $\alpha$ 是学习率（w更新的步长）， $\frac{\partial E}{\partial w}$ 表示梯度，w 更新的方向。
损失函数：网络的输出与真实标签的误差项，训练过程就是不断的最小化误差项 E，得到最优的网络参数【w，b】，
$\sum_x||f_w(x)-y||^2 \tag2$
前向传播：上述公式（2）中的 $f_w(x)$ 表示前向计算过程，x是输入，y是标签；
反向传播：对上述（2）进行链式求导，可以得到全部【w，b】的偏导数（梯度），然后应用公式（1）更新参数【w，b】；
反向传播
关于反向传播的具体推导过程，可以参考链接：https://zhuanlan.zhihu.com/p/40378224，作者写的非常详细，【通俗易懂】，要仔细看看。
分布式特定名称

group：进程组，一般就需要一个默认的，通常使用较少；
world size：全局进程总数量；
rank：进程的ID，0为主进程，一个进程可以控制多个GPU；
local_rank：某个节点上的进程id，隐式参数，torch自动分配；
local_word_size：单个节点上的进程数量，通常很少用到；

图示上述概念，官方参考链接：DDP基本概念及理论，其它参考链接：分布式相关概念（详细）

官方图例解析：假设有两台机器或者节点（NODE），每个机器上有4块GPU，图中有4个进程，即【world_size=4】，进程ID rank=【0,1,2,3】。每一个进程控制【2】块GPU设备。通常情况下，综合考虑设备IO、负载，一般为每一个进程分配【1】块GPU 设备。

在这里插入图片描述

DP

DP (DataParallel) 是单进程，多线程的工作模式，针对单机多卡训练的情况。DP 执行过程中，始终会有一个主GPU来【汇总】和【分发】数据（包括每个GPU的输出，梯度值，网络参数等），通常【0号】GPU被称为主GPU。

GPU数据更新逻辑
1. 主GPU将输入数据（images）平均分发到其它GPU，同时也将模型以及初始的模型参数分发到其它GPU，因此初始的网络参数是一样的，输入图像数据是不同的；
2. 各个GPU独立进行网络前向计算，得到输出；
3. 主GPU将网络输出汇总到主GPU，结合真实的标签值，计算损失，然后对损失求导，记 $\partial loss$ ，然后将该导数分发到其它GPU；
4. 每个GPU得到 $\partial loss$ 后，独立进行反向传播（链式求导），得到所有网络参数的偏导数，也即是梯度；
5. 主GPU将各个GPU的梯度汇总到主GPU，得到平均梯度。根据公式（1）更新网络参数【w，b】，然后将新更新的网络参数分发至其它GPU，这样每个GPU的网络参数依然保持一致；
6. 循环上述过程，至训练结束；
下图的引用链接为：Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups
DP的不足

显卡负载：主GPU的负载较大，通常显存占用比其它显卡大很多。
通信开销：假设有 N 个 GPU，完成一次通信需要时间 t，总共需要花费时间
$T=2(N-1)\cdot t$
性能瓶颈：Global Interpreter Lock (GIL)全局解释器锁，简单来说就是，一个 Python 进程只能利用一个 CPU kernel，即单核多线程并发时，只能执行一个线程。考虑多核，多核多线程可能出现线程颠簸 (thrashing) 造成资源浪费，所以 Python 想要利用多核最好是多进程。

实例展示
YOLO-V5 中的DP使用方式如下，只需添加一行代码即可，如下代码片段所示，

# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
    LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
                   'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
    # 只需添加该行代码即可，其它不必动
    model = torch.nn.DataParallel(model)

DDP

DDP (DistributedDataParallel) 是多进程的工作方式，针对多机多卡（也包含单机多卡）。与 DP 不同的是，不需要全程使用主GPU来汇总和分发数据。

GPU更新逻辑

(1) 各个GPU独立的从硬盘加载输入数据，加载相同的模型；
(2) 各个GPU独立的进行前向传播，损失计算，反向传播；
(3) 每两个相邻的GPU进行梯度的汇总和分发；
(4) 各个GPU进行独立的更新网络参数；
梯度汇总分发策略
对于GPU之间的数据传递，主要采用【The Ring Allreduce】策略。几乎查阅了全网的资源，很少能够解析的比较全面，并且通俗易懂。经过仔细筛选和对比，可以仔细阅读【参考链接，2.，3. 】，详细讲解了DDP的梯度更新策略，并且很容易理解。如果后面有时间，我会详细解析该内容。
通信时间
假设有N个GPU，将传输的梯度分为N份，那么汇总需要传递N-1次，分发需要N-1次。假设传输的数据总量为K，则总的传输时间为，
$2(N-1)\cdot \frac{K}{N}$
注：通过上式可以看到，通信时间几乎与GPU数量无关。

示例展示
YOLO-V5 中的DDP使用方式如下，主要包含两部分：首先，初始化线程组（L12）；然后调用【DDP】API 接口，拷贝模型；代码片段如下所示，

# DDP mode
device = select_device(opt.device, batch_size=opt.batch_size)
if LOCAL_RANK != -1:
    msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
    assert not opt.image_weights, f'--image-weights {msg}'
    assert not opt.evolve, f'--evolve {msg}'
    assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
    assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
    assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
    torch.cuda.set_device(LOCAL_RANK)
    device = torch.device('cuda', LOCAL_RANK)
    dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")

# DDP mode
if cuda and RANK != -1:
    model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)

DP和DDP对比

在DP模式中，总共只有一个进程（受到GIL很强限制），DDP 采用多进程，避免了GIL的限制。
DP 的通信成本随着 GPU 数量线性增长，而 DDP 支持 Ring AllReduce，其通信成本是恒定的，与 GPU 数量无关。
DDP 模式下，每一个GPU的显存占用是基本相同的，负载基本一致；
分布式训练时，基本就使用 DDP模式；

YOLO-V5 实际使用

机器配置：单机，Ubuntu，四张显卡，2080Ti，12G 显存

单机单卡

python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5s.yaml --batch-size 32

DP 训练速度较慢，与单卡训练相比，速度提升有限，不建议使用

python train.py  --batch 32 --data coco.yaml --weights '' --cfg yolov5s.yaml --device 0,1

DDP

nproc_per_node：指定GPU的数量；
device：指定具体使用的 GPU ID；
batch：总共的batch，下面的例子，每一个GPU的输入数据数量为【 64/2=32】；

python -m torch.distributed.run --nproc_per_node 2 train.py 
								--batch 64 
								--data coco.yaml 
								--weights ''
								--cfg yolov5s.yaml 
								--device 0,1