Pytorch分布式训练(单机多卡)

参考博客:https://blog.csdn.net/qq_42255269/article/details/123427094

参考博客:https://blog.csdn.net/xiezongsheng1990/article/details/108713405

我个人是因为要训练一个需要大一点的batch_size的模型,所以用的是torch.distributed这种方式,只实践了这一种方式,但是为了完整性我把第一种torch.nn.DataParallel按照理论也整理过来。

两种方法的优缺点

1 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

2 torch.distributed

优点: 避免了nn.DataParallel的主要缺点,数据不会再分发到主卡上,所以所有卡的显存占用很均匀

缺点: 不友好,调代码需要点精力

1 nn.DataParallel

主要的修改就是用nn.DataParallel处理一下你的model

model = nn.DataParallel(model.cuda(), device_ids=local_rank, output_device=local_rank)

这个很简单,就直接上个例子,根据这个例子去改你的代码就好,主要就是注意对model的修改

注意model要放在主GPU上:model.to(device)

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

gpus = [0, 1, 2, 3]

device = torch.device("cuda", 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()

2 torch.distributed

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

torch.distributed.launch 的启动方式是老的了,现在官方建议用torchrun

2.1 torch.distributed.launch

使用这种方式必须用一下命令运行训练代码:

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

参数说明:

CUDA_VISIBLE_DEVICES 填写要使用的gpu
python -m torch.distributed.launch 是说明用torch.distributed.launch来进行训练
--nproc_per_node=4 这个是每个node你要用的设备数,一个node指一个电脑,一般我们单机多卡时就只有一个node,所以不指定,这里的个数就是你CUDA_VISIBLE_DEVICES的个数

nproc_per_node参数的官方参考链接:https://pytorch.org/docs/stable/distributed.html?highlight=torch%20distributed%20launch#module-torch.distributed.launch

一般我们还需要引入当前项目的环境变量,要不然就容易导致from data.utils import xxx这类代码无法正常找到代码,报错no module data。

一般是在当前shell会话临时导入当前项目目标:

export PYTHONPATH=$PYTHONPATH:/home/python/project_path

另外由于torch.distributed.launch需要用cuda,所以要把cuda导入环境目标,这个可以导入到.bashrc中或着也临时导入:

1 导入到.bashrc

向.bashrc中写入:

export PATH="/usr/local/cuda-11.1/bin:$PATH"

2 临时导入:

export PATH="/usr/local/cuda-11.1/bin:$PATH"

整体你的启动脚本应该大致如下:

train.sh

export PATH="/usr/local/cuda-11.1/bin:$PATH"
export PYTHONPATH=$PYTHONPATH:/home/xxx/python/project_path
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py

为了实现后台训练你可能还需要这个脚本:

train_backup.sh

(bash /home/xxx/python/project_path/train.sh >/dev/null 2>&1 &)

训练代码应该大致如下:

# 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)
# 有sampler的时候shuffle不能为True
train_loader = DataLoader(train_set, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=False, sampler=train_sampler)
#(4)使用 DistributedDataParallel 包装模型
device = torch.device("cuda", local_rank)
model = ...
model = torch.nn.parallel.DistributedDataParallel(model.to(device), device_ids=[local_rank], output_device=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()

2.2 torchrun

torchrun官方参考链接:https://pytorch.org/docs/stable/elastic/run.html#launcher-api

CUDA_VISIBLE_DEVICES=0,2 torchrun --standalone --nnodes=1 --nproc_per_node=2 /media/zxr409/data1/lsg/python_data/SiamRPNPP_Change/train/SiamRPNPP_Transformer_Restnet50_123_Train/train_multi_gpu.py

参数说明:

CUDA_VISIBLE_DEVICES 填写要使用的gpu
--standalone --nnodes=1 说明是只用了一个node(电脑)
--nproc_per_node=2 说明是用这台电脑上的两个设备,这个设备的序号会它自己再编码,比如这里你的CUDA_VISIBLE_DEVICES说要用0,2,等运行的时候它会在程序里认为两个设备的序号是0,1,所以在to(device)的时候,device按照你的gpu个数指定即可。不需要关注真实的设备编号。

train.sh

export PATH="/usr/local/cuda-11.1/bin:$PATH"
export PYTHONPATH=$PYTHONPATH:/home/xxx/python/project_path
CUDA_VISIBLE_DEVICES=0,2 torchrun --standalone --nnodes=1 --nproc_per_node=2 main.py

训练代码应该大致如下:

# main.py
import os
import torch
import argparse
import torch.distributed as dist
#(1)local_rank
local_rank = int(os.environ["LOCAL_RANK"])
#(2)使用 init_process_group 设置GPU 之间通信使用的后端和端口:
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl", init_method='env://')
#(3)使用 DistributedSampler 对数据集进行划分:
train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
# 有sampler的时候shuffle不能为True
train_loader = DataLoader(train_set, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=False, sampler=train_sampler)
#(4)使用 DistributedDataParallel 包装模型
device = torch.device("cuda", local_rank)
model = ...
model = torch.nn.parallel.DistributedDataParallel(model.to(device), device_ids=[local_rank], output_device=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()
文章目录