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