pytorch 单机多卡的正确打开方式

pytorch 使用单机多卡，大体上有两种方式

• 简单方便的 torch.nn.DataParallel(很 low，但是真的很简单很友好)
• 使用 torch.distributed 加速并行训练(推荐，但是不友好)

首先讲一下这两种方式分别的优缺点

• nn.DataParallel
  优点：就是简单
  缺点就是：所有的数据要先load到主GPU上，然后再分发给每个GPU去train，注意这时候主GPU的显存占用很大，你想提升batch_size,那你的主GPU就会限制你的batch_size，所以其实多卡提升速度的效果很有限
  注意： 模型是会被copy到每一张卡上的，而且对于每一个BATCH的数据，你设置的batch_size会被分成几个部分，分发给每一张卡，意味着，batch_size最好是卡的数量n的倍数，比如batch_size=6,而你有n=4张卡，那你实际上代码跑起来只能用3张卡，因为6整除3
• torch.distributed
  优点： 避免了nn.DataParallel的主要缺点，数据不会再分发到主卡上，所以所有卡的显存占用很均匀
  缺点： 不友好，调代码需要点精力，有很多需要注意的问题，我后面会列出

接下来展示如何使用两种方法以及相关注意事项

一、torch.nn.DataParallel

主要的修改就是用nn.DataParallel处理一下你的model
model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

这个很简单，就直接上个例子，根据这个例子去改你的代码就好，主要就是注意对model的修改
注意model要放在主GPU上：model.to(device)

# main.py
import torch
import torch.distributed as dist

gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))

train_dataset = ...
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0]) #注意model要放在主GPU上

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()

二、torch.distributed加速

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，只需要编写一份代码，torch 就会自动将其分配给多个进程，分别在多个 GPU 上运行。

要想把大象装冰箱，总共分四步！！

（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()
print(args.local_rank)

（2）使用 init_process_group 设置GPU 之间通信使用的后端和端口：

dist.init_process_group(backend='nccl')

（3）使用 DistributedSampler 对数据集进行划分：

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 = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

• 举个栗子，参照这个例子去设置你的代码结构

# 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()

然后，使用以下指令，执行你的主脚本，其中--nproc_per_node=4表示你的单个节点的GPU数量：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

问题来了！！

你可能会在完成代码之后遇到各种问题，我这里列举一些要注意的点，去避坑
如果你遇到的莫名奇妙报错的问题，尝试这样去修改你的代码

• device 的设置
  你需要设置一个device参数，用来给你的数据加载到GPU上，由于你的数据会在不同线程中被加载到不同的GPU上，你需要传给他们一个参数device,用于a.to(device)的操作（a是一个tensor）
  device如下设置

device = torch.device("cuda", args.local_rank)

你也可以通过设置当前cuda，使用a.cuda()把张量放到GPU上，但是不推荐，可能会有一些问题

torch.cuda.set_device(args.local_rank)

• find_unused_parameters=True
  这个是为了解决你的模型中定义了一些在forward函数中没有用到的网络层，会被视为“unused_layer”,这会引发错误，所以你在使用 DistributedDataParallel 包装模型的时候，传一个find_unused_parameters=True的参数来避免这个问题，如下：

encoder=nn.parallel.DistributedDataParallel(encoder, device_ids=[args.local_rank],find_unused_parameters=True)

• num_workers
  很好理解，尽量不要给你的DataLoader设置numworkers参数，可能会有一些问题（不要太强迫症）
• shuffle=False
  你的DataLoader不要设置shuffle=True

valid_loader = torch.utils.data.DataLoader(
        part_valid_set, batch_size=BATCH, shuffle=False, num_workers=num_workers,sampler=valid_sampler)