DistributedDataParallel多GPU分布式训练全过程总结 跟着做90%成功

Posted 2022-03-25 BRiAq

DistributedDataParallel分布式训练全教程

• • 什么是DistributedDataParallel
  • 为什么是DistributedDataParallel
  • 要不要装别的库
  • 教程
  • • 一、初始化
    • 二、加一个解析参数
    • 三、数据集
    • 四、设定device
    • 五、模型加载到多gpu
    • 六、数据加载到gpu
    • 七、最后一句命令
  • 参考

什么是DistributedDataParallel

一种基于pytorch框架的分布式训练工具（也就是让你用多GPU来训练）

为什么是DistributedDataParallel

nn.DataParallel已经被淘汰了，连官方文档都建议替换成DistributedDataParallel，然后我自己用也问题连天，各种小bug，调了两天bug硬是没跑下来，DistributedDataParallel一下午就调通了（这不是正在总结教程），然后三块gpu的使用还很平均，可见好用

要不要装别的库

torch就够

教程

一、初始化

这句必须有，跟着写就完了

torch.distributed.init_process_group(backend="nccl")

以下两句最好有

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"  # 有几块GPU写多少

二、加一个解析参数

在运行的时候，DistributedDataParallel会往你的程序中加入一个参数local_rank，所以要现在你的代码中解析这个参数，如：

parser.add_argument("--local_rank", type=int, default=1, help="number of cpu threads to use during batch generation")

不加会报错，提示多了local_rank参数
不知道什么是解析参数的赶紧去百度，都已经做到分布式训练了肯定要补课

三、数据集

简单的三步走战略：

1. 得到你的DataSets类的数据集，我们暂时叫train_dataset
2. 获得它的sampler：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

3.将dataset和sampler同时装进dataloader:

train_loader = torch.utils.data.DataLoader(
	train_dataset,
	batch_size = args.batch_size,
	shuffle = False,
	num_workers = args.num_workers,
	pin_memory = True,
	sampler = train_sampler)
	
# pin_memory=true 能加快内存的Tensor转义到GPU的显存的速度，
# 前提是你的内存要足够(一般为使用显卡显存x2)，如果不够，建议关闭pin_memory(锁页内存)选项

很容易吧？这之后dataloader就一样的用了
4. 同样的方式处理val, test等

四、设定device

local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
global device
device = torch.device("cuda", local_rank)

我没用arg.local_rank，新定义了一个local_rank变量，是因为我更信任distributed.get_rank()这个函数
这里用torch.device来写，并且加了global，是因为后面模型和数据都要用到这个device，不会出错

五、模型加载到多gpu

model.to(device)  # 这句不能少，最好不要用model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)  # 这句加载到多GPU上

六、数据加载到gpu

数据.to(device)

明白吗，就用这样的写法

七、最后一句命令

最终，在命令行运行你的py文件（假设是train.py）

python -m torch.distributed.launch --nnodes=1 --nproc_per_node=3 train.py

第一：nnodes是你的机器数，只用一台电脑的话，就是1
第二：nproc_per_node是你的gpu数，有几个写几
第三：就是要通过命令行来跑你的程序，有人问那我想单步调试呢，我用pycharm行不行？答案是会报错。如果有知道怎么可以实现单步调试的也欢迎在下面留言提供方法

运行的过程中，nvidia-smi可以明显看到多张卡分配的均匀的占用，训练速度明显变快，nproc_per_node调少那么后几个gpu也不会用

参考

在调通过程中，感谢一些大佬的文章参考：
1.https://zhuanlan.zhihu.com/p/95700549
2.https://zhuanlan.zhihu.com/p/206467852
3.DistributedDataParallel官方文档
4.数据被放进不同gpu报错的问题解决