pytorch单机多卡的正确打开方式 以及可能会遇到的问题和相应的解决方法

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)