Pytorch ddp多卡训练踩坑(1)

多卡训练的时候发现一个问题：cuda() 方法会无视环境变量CUDA_VISIBLE_DEVICES。示例如下：

CUDA_VISIBLE_DEVICES="3,4,5,6" python -m torch.distributed.launch --nproc_per_node=${N_GPUS} train.py --config_option <your_config_here>

这边指定了可见显卡为3-6号，然而如果在代码里出现cuda()，模型/数据还是会load到0卡(default)。

# load to GPU_0
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], broadcast_buffers=False, find_unused_parameters=False)
sample.cuda()

解决方案1，初始化ddp的时候设置默认cuda设备，然而这个方法对我没用。

torch.cuda.set_device(args.local_rank)

解决方案2(final)，显式指定device/GPU_id，以0号进程为例，代码如下：

# load to GPU_3
device = torch.device("cuda:0")
model.to(device)
model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], broadcast_buffers=False, find_unused_parameters=False)
sample.to(model.device)

这里的cuda:0其实是CUDA_VISIBLE_DEVICES的第一张卡，即3号卡。当然，上述只是举例，旨在说明CUDA_VISIBLE_DEVICES的作用。

正确写法是model.to(local_rank)，不同进程直接load模型到对应的GPU上。(之前解法1+cuda()也是ok的，不知道为什么在新的节点上不行~)