pytorch 使用单机多卡,大体上有两种方式
首先讲一下这两种方式分别的优缺点
batch_size
,那你的主GPU就会限制你的batch_size
,所以其实多卡提升速度的效果很有限batch_size
会被分成几个部分,分发给每一张卡,意味着,batch_size
最好是卡的数量n
的倍数,比如batch_size=6
,而你有n=4
张卡,那你实际上代码跑起来只能用3
张卡,因为6整除3接下来展示如何使用两种方法以及相关注意事项
主要的修改就是用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()
与 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)
forward
函数中没有用到的网络层,会被视为“unused_layer”,这会引发错误,所以你在使用 DistributedDataParallel 包装模型的时候,传一个find_unused_parameters=True
的参数来避免这个问题,如下:encoder=nn.parallel.DistributedDataParallel(encoder, device_ids=[args.local_rank],find_unused_parameters=True)
DataLoader
设置numworkers
参数,可能会有一些问题(不要太强迫症)DataLoader
不要设置shuffle=True
valid_loader = torch.utils.data.DataLoader(
part_valid_set, batch_size=BATCH, shuffle=False, num_workers=num_workers,sampler=valid_sampler)