DistributedDataParallel分布式训练全教程
什么是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参数
不知道什么是解析参数的赶紧去百度,都已经做到分布式训练了肯定要补课
三、数据集
简单的三步走战略:
- 得到你的DataSets类的数据集,我们暂时叫train_dataset
- 获得它的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报错的问题解决
部分踩坑及解决
–22.9.1更新 模型存储的注意事项!–
采用DistributedDataParallel并行方式训练,在test时,如果遇到以下的报错:
Traceback (most recent call last):
File “/root/PycharmProjects/test.py”, line 8, in
model_dict = torch.load(“0.pth”)
File “torch/serialization.py”, line 529, in load
return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
File “torch/serialization.py”, line 709, in _legacy_load
deserialized_objects[key]._set_from_file(f, offset, f_should_read_directly)
RuntimeError: storage has wrong size: expected -4916312287391674656 got 24
这是因为模型存储时被多个进程同时写入,因而产生了写入结果存在冗余的情况!
解决方法:存储模型时通过local_rank限制单进程写入
def save_checkpoint(epoch, model, best_top5, optimizer,
is_best=False,
filename='checkpoint.pth.tar'):
state = {
'epoch': epoch+1, 'state_dict': model.module.state_dict(),
'best_top5': best_top5, 'optimizer' : optimizer.state_dict(),
}
torch.save(state, filename)
if args.local_rank == 0:
if is_best: save_checkpoint(epoch, model, best_top5, optimizer, is_best=True, filename='model_best.pth.tar')
注意,在DDP中多了一个module,所以保存的是model.module.state_dict()
参考文章