单机多卡训练神经网络

DDP（DistributedDataParallel）

DDP优于DP，这里不做赘述，请读者自己搜索吧。

如何迅速集成？

直接上代码：
第一步：初始化

dist.init_process_group(backend='nccl')
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

第二步：将模型并行化

if torch.cuda.device_count() > 1:
	model = torch.nn.parallel.DistributedDataParallel(model .to(device),find_unused_parameters=True,broadcast_buffers=False,device_ids=[local_rank],output_device=local_rank)

第三步：将数据放在GPU上：

inputs = inputs.to(device)
gt= gt.to(device)

笔者遇到的问题

RuntimeError: CUDA out of memory

解决方案

原来的代码

torch.nn.parallel.DistributedDataParallel(model .to(device),find_unused_parameters=True,broadcast_buffers=False,device_ids=[0，1，2，3])

修改后的代码

model = torch.nn.parallel.DistributedDataParallel(model .to(device),find_unused_parameters=True,broadcast_buffers=False,device_ids=[local_rank],output_device=local_rank)

貌似是放不下导致的问题，但是单卡训练是可以的，最后这样解决了问题。