Pytorch：多块GPU分布式|并行训练

分布式与并行训练的区别

• 分布式：
  多台服务器上的多个GPU，分布式涉及了服务器之间的通信，因此比较复杂，PyTorch封装了相应的接口，可以用几句简单的代码实现分布式训练。
• 并行：
  一台服务器上的多个GPU

多GPU训练

可以分为model parallel（模型并行）和data parallel（数据并行）

• model parallel
  由于模型太大了，单块GPU跑不起来，因此需要将一个模型分到不同的GPU上。没有加速训练的作用。
• data parallel
  将数据分到不同的GPU上，可以加速训练

多GPU并行训练时会涉及到多GPU之间的通信问题，因此训练时间不是线性递减的。

在Pytorch实现多GPU训练

Pytorch提供了两种多GPU的训练，DataParallel和DistributedDataParallel

• DataParallel:
  早期方法，仅能实现单机多卡，只有一个进程，master节点相当于参数服务器，向其他服务器广播参数，梯度正向传播后，各GPU将梯度的数值传回到master节点上，master节点对搜集来的梯度进行计算后更新参数，再广播出去更新后的参数。

  这种方式使得master节点计算任务重，阻塞网络，训练速度不是很高

  相比于单机单卡仅需要增加一行代码

  #创建模型
  model = models.resnet101()
  #模型拷贝
  model = torch.nn.DataParallel(model.cuda(), device_ids = [0,1,2,3])
  

• DistributedDataParallel:
  推荐方法，可以实现单机多卡、多机多卡，而且速度更快

具体实现步骤

以下仅实现单机多卡

1 启动方式

torch.distributed.launch

2 传入参数

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
parser.add_argument("--gpu_id", type=str, default='0,1,2,3,4,5', help='path log files')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = 'PCI_BUS_ID'
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_id

3 初始化进程

torch.distributed.init_process_group("nccl"，world_size=n_gpu，rank=args.local_rank) # 第一参数nccl为GPU通信方式, world_size为当前机器GPU个数,rank为当前进程在哪个PGU上

4 设置进程使用的卡

torch.cuda.set_device(args.local_rank)

5 对模型进行包裹

model=torch.nn.DistributedDataParallel(model.cuda(args.local_rank), device_ids=[args.local_rank]),#这里device_ids传入一张卡即可，因为是多进程多卡，一个进程一个卡

6 将数据分配到不同的GPU上

train_sampler = torch.util.data.distributed.DistributedSampler(train_dataset) # train_dataset为Dataset()

将train_sampler传入到DataLoader中,不需要传入shuffle=True,因为shuffle和sampler互斥 data_dataloader = DataLoader(…, sampler=train_sampler)

整体代码示例

# main.py
import torch
import argparse
import torch.distributed as dist
#（1）要使用`torch.distributed`，你需要在你的`main.py(也就是你的主py脚本)`中的主函数中加入一个**参数接口：`--local_rank`**
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()
#（2）使用 init_process_group 设置GPU 之间通信使用的后端和端口：
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)
#（3）使用 DistributedSampler 对数据集进行划分：
train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)
#（4）使用 DistributedDataParallel 包装模型
model = ...
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

运行命令

python -m torch.distributed.launch main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset Grocery_and_Gourmet_Food

参考文章

pytorch多GPU分布式训练代码编写
Pytorch中多GPU并行计算教程