Pytorch 分布式训练

Posted 2022-08-30 洪流之源

Yolov5支持多机多卡分布式训练，参与分布式训练的机器用Node表述，没有使用Host，
原因是Node不限定是物理机器，还有可能是容器例如docker，简单理解就是一个Node节点就是一台机器，
Node又分为Master Node、Slave Node，Master Node只有一个，Slave Node可以有多个
假定现在有两台机器参与分布式训练，每台机器有4张显卡，可以用下面的命令分别启动Master Node、Slave Node的训练:

python -m torch.distributed.launch --nnodes 2 --nproc_per_node 4 --node_rank 0 --master_addr "192.168.1.2" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights 'yolov5s.pt'

python -m torch.distributed.launch --nnodes 2 --nproc_per_node 4 --node_rank 1 --master_addr "192.168.1.2" --master_port 1234 train.py --batch 64 --data coco.yaml --cfg yolov5s.yaml --weights 'yolov5s.pt'

上述的命令中
--nnodes表示一共多少台机器参与分布式训练，也就是有几个Node，只有2台机器所以设置为2
--nproc_per_node表示每台机器有多少张显卡，每台机器有4张显卡所以设置为4
--node_rank表示当前机器的序号，一般设置为0的作为Master Node
--master_add表示Master Node的ip地址
--master_add表示Master Node的端口号
一般情况下，当两台机器都运行完命令，训练就开始了，否则master处于等待状态

如下所示，训练开始后Node节点的分配情况：
Node 0
    Process0 [Global Rank=0, Local Rank=0] -> GPU 0
    Process1 [Global Rank=1, Local Rank=1] -> GPU 1
    Process2 [Global Rank=2, Local Rank=2] -> GPU 2
    Process3 [Global Rank=3, Local Rank=3] -> GPU 3
Node 1
    Process4 [Global Rank=4, Local Rank=0] -> GPU 0
    Process5 [Global Rank=5, Local Rank=1] -> GPU 1
    Process6 [Global Rank=6, Local Rank=2] -> GPU 2
    Process7 [Global Rank=7, Local Rank=3] -> GPU 3

每一张显卡被分配一个进程，从Process0 ~ Process7，Global Rank表示在整个分布式训练任务中的分布式进程编号，从Global Rank = 0 ~ Global Rank = 7
Local Rank表示在某个Node内的编号，在Node 0中Local Rank = 0 ~ Local Rank = 3
另外，如果只写rank前面没有global、local等字段一般指代Global Rank
world_size表示全局进程数量，也就是分布式进程的数量，在上述的配置中world_size = 8，如果一共有3个node（nnodes=3），每个node包含8个GPU，设置nproc_per_node=4，
world_size就是3 * 4 = 12，为什么不是3 * 8 = 24呢？因为每个node虽然有8个GPU，但是命令设置只使用其中4个（nproc_per_node=4），有而不使用是不算数的。

参考：Pytorch中多GPU并行计算教程_太阳花的小绿豆的博客-CSDN博客_pytorch 多gpu