单机多卡DDP训练的一般流程。一些代码可能仅支持单卡train，这时就需要自己手动修改代码，提高训练效率。
大致流程
导入所需的库：
import torch
import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler
设置local_rank argparse参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()
初始化DDP，设置进程通信的后端为nccl，同步所有进程，并获取当前进程组内所有进程数
dist.init_process_group(backend='nccl')
# 同步所有进程
dist.barrier()
# 获取当前进程组内所有进程数
world_size = dist.get_world_size()
将数据集分给每一个进程，避免进程间数据重复
train_sampler = DistributedSampler(train_dataset)
valid_sampler = DistributedSampler(valid_dataset)
将sampler传入dataloader
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)
valid_loader = DataLoader(valid_dataset, sampler=valid_sampler, batch_size=256, pin_memory=False, prefetch_factor=2, num_workers=4)
假设模型已经load，可以直接使用.to(device)将模型分配至CUDA
device = torch.device("cuda", args.local_rank)
model = model.to(device)
使用SyncBN优化各自进程中数据集的BN计算
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
使用DistributedDataParallel将模型包装起来
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
在每次迭代前，通过set_epoch设置random seed，使得每次运行时shuffle每个epoch的数据，并让每个GPU拿到的数据不同
train_sampler.set_epoch(epoch)
valid_sampler.set_epoch(epoch)
启动分布式训练
# 单机多卡
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_example.py
# 单机多卡，但对一个代码执行不同实验，需要指定端口
$ CUDA_VISIBLE_DEVICES=2,3 python -m torch.distributed.launch --nproc_per_node=2 --master_port 9999 ddp_example.py
避免重复输出
参考maskrcnn项目，利用is_main_process()函数判断是否属于主进程。如果是主进程，再对信息进行输出。
对于tqdm进度条， 可以加入disable参数进行控制，使得仅在主进程时输出进度条，例如
with tqdm(enumerate(dataloader), total=len(dataloader), ncols=120, disable=not is_main_process()) as t:
保证tensor运算在同一设备上
有的代码在单卡改多卡时，会报错： **RuntimeError: Expected all tensors to be on the same device, but found at least two devices,**
这是因为原来的单卡代码中，在为tensor指定gpu时，使用的是.cuda()，但为model分配gpu时，使用的是model.to(device)，从而导致模型输出的tensor和之前的tensor不在同一张显卡上，二者运算时报错。
最好将报错处的所有tensor的.cuda()改成.to(device)，统一运算设备。
文档信息
- 本文作者：焦逸凡
- 本文链接：https://ailovejinx.github.io/2023/11/01/blog-ddp/
- 版权声明：自由转载-非商用-非衍生-保持署名（创意共享3.0许可证）