PyTorch中DistributedDataParallel 使用笔记
1. 为何使用DistributedDataParallel
在PyTorch中的多GPU训练一般有2种DataParallel和DistributedDataParallel，DataParallel是最简单的的单机多卡实现，但是它使用多线程模型，并不能够在多机多卡的环境下使用，所以本文将介绍DistributedDataParallel，DDP 基于使用多进程而不是使用多线程的DP，并且存在GIL争用问题，并且可以扩充到多机多卡的环境，所以它是分布式多GPU训练的首选。
2. 基本概念
在使用DistributedDataParallel时有一些概率必须掌握
| 多机多卡 | 含义 |
|---|---|
| world_size | 代表有几台机器，可以理解为几台服务器 |
| rank | 第几台机器，即第几个服务器 |
| local_rank | 某台机器中的第几块GPU |
| 单机多卡 | 含义 |
|---|---|
| world_size | 代表机器一共有几块GPU |
| rank | 第几块GPU |
| local_rank | 第几块GPU，与rank相同 |
3. 使用方法
3.1. 修改主函数
在运行的时候，DistributedDataParallel会往你的程序中加入一个参数local_rank，所以要现在你的代码中解析这个参数，如：
parser.add_argument("--local_rank", type=int, default=1, help="number of cpu threads to use during batch generation")3.2. 初始化
torch.distributed.init_process_group(backend="nccl") os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2" # 有几块GPU写多少3.3. 设定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，不会出错
3.4. 模型加载到多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上3.5. 数据加载到gpu
数据.to(device)3.6. 启动
torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:12345 train_cylinder_asym.py参考文献
Pytorch并行计算(二): DistributedDataParallel介绍_dist.barrier_harry_tea的博客-CSDN博客
DistributedDataParallel多GPU分布式训练全过程总结 跟着做90%成功_BRiAq的博客-CSDN博客
百度安全验证