pytorch多GPU并行以及注意事项

pytorch多GPU并行方式

1-参数解析

使用argparse进行参数解析，核心参数如下：

import argparse
parser.add_argument("--local_rank", default=0, type=int)
parser.add_argument("--ngpu", default=6, type=int)
args.gpu = args.local_rank
torch.cuda.set_device(args.gpu)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()

2-使用数据并行

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device= args.local_rank)
# dataset为torch.utils.data.Dataset的实例
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
train_loader = torch.utils.data.DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=(train_sampler is None),
                                           num_workers=4, pin_memory=True, sampler=train_sampler, drop_last=True)

3-保存权重

if args.local_rank == 0:
	torch.save(model.module.state_dict(), 'res.pth')

4-运行程序

pytorch多GPU运行时，有特定的运行命令：

# 6为GPU的总数
python -m torch.distributed.launch --nproc_per_node=6 train.py

5-注意事项

DistributedDataParallel会将model进行封装，容易产生如下两个问题：

• 在编写GAN模型时，如果要用模型里面的层，需要用model.module.x_layer而不是直接用model.x_layer
• 想在训练时保存权重，也需要在加载权重时修改权重key的名称，比如将名称由stage1.layer1改为module.stage1.layer1，否则就会出现无法加载的情况（如下图）