torch distributed 多GPU训练笔记

 

DistributedDataParallel (既可单机多卡又可多机多卡)

先奉上官网nn.DistributedDataParallel(model)链接

https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel​pytorch.org

1. CLASS torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None,

2. dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25,

3. find_unused_parameters=False, check_reduction=False)

• 参数定义

module是要放到多卡训练的模型;

device_ids数据类型是一个列表, 表示可用的gpu卡号;

output_devices数据类型也是列表,表示模型输出结果存放的卡号(如果不指定的话,默认放在0卡,这也是为什么多gpu训练并不是负载均衡的,一般0卡会占用的多,这里还涉及到一个小知识点——如果程序开始加os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3", 那么0卡(逻辑卡号)指的是2卡(物理卡号))。

dim