pytorch分布式训练代码编写

一、单机单卡

1. 模型拷贝（原地操作）

model.cuda()

2. 数据拷贝（赋值操作）

data = data.cuda()

3. 判断GPU是否可用

torch.cuda.is_avaliable()

4. 模型保存与加载

torch.save

torch.load(file.pt,map_location=torch.device(cuda))

二、单机多卡

1. 检测GPU数目

torch.cuda.device_count()

或者使用命令行CUDA_VISIBLE_DEVICES=" "来限制GPU的使用

2. torch.nn.DataParallel（不推荐）
   该API已经被淘汰，原因是不支持多进程、不支持多机情况、不支持模型并行。

3. torch.nn.parallel.DistributedDataParallel（推荐）
   初始化进程组：

torch.distributed.init_process_group(backend='nccl', world_size=n_gpus, rank=args.local_rank, init_method='...')

'nccl’是GPU之间的通信方式，world_size指的是当前机器上用几张GPU，rank用来指定当前进程在哪个GPU上

torch.cuda.set_device(args.local_rank)

该语句相当于作用CUDA_VISIBLE_DEVICES环境变量。

准备工作完成后，用API包裹模型：

model = torch.nn.parallel.DistributedDataParallel(
			model.cuda(args.local_rank), 
			device_ids=[args.local_rank]
    )

指定sampler：

train_sampler = DistributedSampler(train_dataset)

train_dataloader = DataLoader(...,sampler=train_sampler,...)

data = data.cuda(args.local_rank)

执行命令：

python -m torch.distributed.launch --nproc_per_node=n_gpus train.py

4. 模型保存
   torch.save在local_rank=0的位置进行保存，注意调用model.module.state_dict()
torch.load注意map_location