PyTorch多卡分布式训练DistributedDataParallel 使用方法
PyTorch多卡分布式训练DistributedDataParallel 使用方法
目录
PyTorch多卡分布式训练DistributedDataParallel 使用方法
1.DP模式和DP模式
(1)单进程多GPU训练模式:DP模式
(2)多进程多GPU训练模式：DDP模式
2.Pytorch分布式训练方法
3.Pytorch-Base-Trainer(PBT)分布式训练工具
(1) 工具介绍
(2) 安装
(3)使用方法
4.Example: 构建自己的分类Pipeline
5.可视化
【尊重原则，转载请注明出处】：https://panjinquan.blog.csdn.net/article/details/122702287
1.DP模式和DP模式
Pytorch多卡训练有两种方式，一种是单进程多GPU训练模式(single process multi-gpus)，另一种的多进程多卡模式(multi-processes multi-gpus)
(1)单进程多GPU训练模式:DP模式
Pytorch通过nn.DataParallel可实现多卡训练模型（简称DP模式），这是single process multi-gpus 的多卡并行机制，这种并行模式下并行的多卡都是由一个进程进行控制，其缺点有：
- 尽管 DataLoader 可以指定 num_worker，增加负责加载数据的线程数量，但是线程的资源受限于父进程，且由于python的GIL机制，不能利用好多核的并行优势
- 模型在 gpu 群组中进行初始化与广播过程依赖单一进程的串行操作
- DP模式相当于将多个GPU卡合并为一个卡进行训练
尽管DataParallel更易于使用（只需简单包装单GPU模型），但由于使用一个进程来计算模型权重，然后在每个批处理期间将分发到每个GPU，因此通信很快成为一个瓶颈，GPU利用率通常很低。而且，nn.DataParallel要求所有的GPU都在同一个节点上（不支持分布式）。
(2)多进程多GPU训练模式：DDP模式
Pytorch通过nn.parallel.DistributedDataParallel可实现多进程多卡训练模型（也称DDP模式），这种多卡并行机制的特点/优势有：
- 一个进程一个GPU(当然可以让每个进程控制多个GPU，但这显然比每个进程有一个GPU要慢)
- 充分利用多核并行的优势加载数据
- 模型在 gpu 群组中进行初始化的过程由各自的进程负责调度
- 代码可以无缝切换单机多卡与多机多卡训练，因为此时单机单卡成为了单机多卡/多机多卡并行下的一个特例
- GPU可以都在同一个节点上，也可以分布在多个节点上。每个进程都执行相同的任务，并且每个进程都与所有其他进程通信。进程或者说GPU之间只传递梯度，这样网络通信就不再是瓶颈。
在训练过程中，每个进程从磁盘加载batch数据，并将它们传递到其GPU。每一个GPU都有自己的前向过程，然后梯度在各个GPUs间进行All-Reduce。每一层的梯度不依赖于前一层，所以梯度的All-Reduce和后向过程同时计算，以进一步缓解网络瓶颈。在后向过程的最后，每个节点都得到了平均梯度，这样模型参数保持同步。
这就要求多个进程，甚至多个节点上的多个进程实现同步并通信。Pytorch通过distributed.init_process_group函数来实现这一点。他需要知道进程0位置以便所有进程都可以同步，以及预期的进程总数。每个进程都需要知道进程总数及其在进程中的顺序，以及使用哪个GPU。通常将进程总数称为world_size。
Pytorch提供了nn.utils.data.DistributedSampler来为各个进程切分数据，以保证训练数据不重叠。
nn.DataParallel和nn.distributedataparallel的主要差异可以总结为以下几点：
DistributedDataParallel支持模型并行，而DataParallel并不支持，这意味如果模型太大单卡显存不足时只能使用前者；DataParallel是单进程多线程的，只用于单机情况，而DistributedDataParallel是多进程的，适用于单机和多机情况，真正实现分布式训练；DistributedDataParallel的训练更高效，因为每个进程都是独立的Python解释器，避免GIL问题，而且通信成本低其训练速度更快，基本上DataParallel已经被弃用；- 必须要说明的是
DistributedDataParallel中每个进程都有独立的优化器，执行自己的更新过程，但是梯度通过通信传递到每个进程，所有执行的内容是相同的；
除了PyTorch官方实现的分布式训练方案，还有horovod分布式训练工具，不仅支持PyTorch还支持TensorFlow和MXNet框架，实现起来也是比较容易的，速度方面应该不相上下。
参考资料：PyTorch分布式训练简明教程 - 知乎
2.Pytorch分布式训练方法
分布式训练一般分为数据并行和模型并行两种，Pytorch分布式训练的实现步骤可简述如下：
- 首先在nn.DataParallel(即DP模式下)实现多卡加载数据，训练模型并调试成功；这一步是为了保证你的训练流程正常无BUG。然后就可以开始魔改了
- 数据并行(分布式)：
DataLoader的样本采样器(sampler)修改为分布式采样器torch_utils.distributed.DistributedSampler - 模型并行(分布式)：将torch.nn.parallel.DistributedDataParallel 代替torch.nn.DataParallel
- 为了能够使用 DistributedDataParallel 需要先进行进程间通讯环境的初始化，torch.distributed.init_process_group()
- 为了解决并行训练中加载到各个 worker/gpu 中的 sub-mini-batch 之间出现 example overlap 问题，还可以配合 torch.utils.data.distributed.DistributedSampler 进行使用
- 为了让进程与 gpu 进行一一匹配，在程序的开头通过 torch.cuda.set_device 设定目标设备
- （可选）为了让各个 worker/gpu 能有一致的初始值，在程序开头通过 torch.manual_seed 与 torch.cuda.manual_seed 来初始化随机数种子
所以代码结构如下：
# filename: distributed_example.py
# import some module
...
...
parser = argparse.Argument()
parser.add_argument('--init_method', defalut='env://', type=str)
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse()
import os
# Set master information and NIC
# NIC for communication
os.environ['NCCL_SOCKET_IFNAME'] = 'xxxx'
# set master node address
# recommend setting to ib NIC to gain bigger communication bandwidth
os.environ['MASTER_ADDR'] = '192.168.xx.xx'
# set master node port
# **caution**: avoid port conflict
os.environ['MASTER_PORT'] = '1234'
def main():
# step 1
# set random seed
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
# step 2
# set target device
torch.cuda.set_device(args.local_rank)
# step 3
# initialize process group
# use nccl backend to speedup gpu communication
torch.distributed.init_process_group(backend='nccl', init_method=args.init_method)
...
...
# step 4
# set distributed sampler
# the same, you can set distributed sampler for validation set
train_sampler = torch.utils.data.distributed.DistributedSampler(
dataset_train)
train_loader = torch.utils.data.DataLoader(
dataset_train, batch_size=BATCH_SIZE, sampler=train_sampler, pin_memory=PIN_MEMORY,
num_workers=NUM_WORKERS
)
...
...
# step 5
# initialize model
model = resnet50()
model.cuda()
# step 6
# wrap model with distributeddataparallel
# map device with model process, and we bind process n with gpu n(n=0,1,2,3...) by default.
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
...
...
for epoch in range(epochs):
# step 7
# update sampler with epoch
train_sampler.set_epoch(epoch)
# then do whatever you want to do, just like a single device/gpu training program
...
...注意：
由于程序开头采用了 torch.cuda.set_device() 指定了目标的设备，所以后续的操作中如果有涉及要将数据、模型移动到 gpu 中的操作需要改为 model.cuda()、inputs.cuda()，该指令会将对象正确地复制到对应的 cuda 设备中。
如果你执意使用 to 操作，那么注意确保 xx.to('cuda:n') 中的 cuda:n 与目标设备是匹配的。
以多进程方式启动训练脚本
你当然可以以 python distributed_example.py 这样的形式启动训练脚本，不过这样无法触发多进程。pytorch 为多进程启动脚本提供了 launch 工具，所以正确的启动方式是：
python -m torch.distributed.launch --nnodes= --nproc_per_node= --node_rank=\
distributed_example.py --arg1 --arg2 and all other arguments of your trainning script 参数说明：
- nnodes：指定参与计算的节点数量，默认值为1，单机多卡的训练中可以不用指定
- nproc_per_node：指定每个节点中的所要启动的进程数量，由于进程与 gpu 一一对应，所以这里的数值不能大于系统中所能使用的 gpu 数量
- node_rank: 指定当前节点在整个系统中的序号，从 0 开始递增，需要注意的是，在多机多卡训练中 node_rank == 0 的节点表示 master，所以 node_rank == 0 的节点必须是 MASTER_ADDR 所在的节点，否则多卡间的通信无法正确建立连接。
老实说，【从DP模式升级到DDP的方法】看起来简单，步骤也不多，但真正要跑起来还是很多地方需要优化的；
这种多进程训练的方法，每个进程需要分配一个卡进行训练，这就导致你保存模型，打印Log,测试数据都变成复杂了，比如会出现多个进程都会打印相同的Log的问题；一般建议你，定义一个主进程，且在主进程中打印Log,保存模型，测试数据等操作，这样可以避免上述问题了。
那有没有一个简单方法，可以快速实现Pytorch的分布式训练
有的，我今天就介绍一个我自己整合的Pytorch的分布式训练工具：Pytorch-Base-Trainer，基于这套工具，你可以简单配置，即可实现DP或者DDP模式的训练，而无需关注各种进程间通讯，端口设置等这些复杂的过程。
3.Pytorch-Base-Trainer(PBT)分布式训练工具
(1) 工具介绍
考虑到深度学习训练过程都有一套约定成俗的流程，鄙人借鉴Keras开发了一套基础训练库： Pytorch-Base-Trainer(PBT)； 这是一个基于Pytorch开发的基础训练库，支持以下特征：
- 支持多卡训练训练(DP模式)和分布式多卡训练(DDP模式)，参考build_model_parallel
- 支持argparse命令行指定参数，也支持config.yaml配置文件
- 支持最优模型保存ModelCheckpoint
- 支持自定义回调函数Callback
- 支持NNI模型剪枝(L1/L2-Pruner,FPGM-Pruner Slim-Pruner)nni_pruning
- 非常轻便,安装简单
博客介绍：
Pytorch基础训练库Pytorch-Base-Trainer(支持模型剪枝 分布式训练)_pan_jinquan的博客-CSDN博客考虑到深度学习训练过程都有一套约定成俗的流程，鄙人借鉴Keras开发了一套基础训练库： Pytorch-Base-Trainer(PBT)； 这是一个基于Pytorch开发的基础训练库，支持以下特征：https://panjinquan.blog.csdn.net/article/details/122662902GitHub地址：
GitHub - PanJinquan/Pytorch-Base-Trainer: Pytorch分布式训练框架
https://github.com/PanJinquan/Pytorch-Base-Trainer
(2) 安装
- 源码安装
git clone https://github.com/PanJinquan/Pytorch-Base-Trainer
cd Pytorch-Base-Trainer
bash setup.sh #pip install dist/basetrainer-*.*.*.tar.gz- pip安装
pip install basetrainer- 使用NNI 模型剪枝工具，需要安装NNI
# Linux or macOS
python3 -m pip install --upgrade nni
# Windows
python -m pip install --upgrade nni(3)使用方法
basetrainer使用方法可以参考example.py,构建自己的训练器,可通过如下步骤实现：
- step1: 新建一个类
ClassificationTrainer，继承trainer.EngineTrainer - step2: 实现接口
def build_train_loader(self, cfg, **kwargs):
"""定义训练数据"""
raise NotImplementedError("build_train_loader not implemented!")
def build_test_loader(self, cfg, **kwargs):
"""定义测试数据"""
raise NotImplementedError("build_test_loader not implemented!")
def build_model(self, cfg, **kwargs):
"""定于训练模型"""
raise NotImplementedError("build_model not implemented!")
def build_optimizer(self, cfg, **kwargs):
"""定义优化器"""
raise NotImplementedError("build_optimizer not implemented!")
def build_criterion(self, cfg, **kwargs):
"""定义损失函数"""
raise NotImplementedError("build_criterion not implemented!")
def build_callbacks(self, cfg, **kwargs):
"""定义回调函数"""
raise NotImplementedError("build_callbacks not implemented!")
step3: 在初始化中调用build
def __init__(self, cfg):
super(ClassificationTrainer, self).__init__(cfg)
...
self.build(cfg)
...
step4: 实例化ClassificationTrainer,并使用launch启动分布式训练
def main(cfg):
t = ClassificationTrainer(cfg)
return t.run()
if __name__ == "__main__":
parser = get_parser()
args = parser.parse_args()
cfg = setup_config.parser_config(args)
launch(main,
num_gpus_per_machine=len(cfg.gpu_id),
dist_url="tcp://127.0.0.1:28661",
num_machines=1,
machine_rank=0,
distributed=cfg.distributed,
args=(cfg,))4.Example: 构建自己的分类Pipeline
basetrainer使用方法可以参考example.py
# 单进程多卡训练
python example.py --gpu_id 0 1 # 使用命令行参数
python example.py --config_file configs/config.yaml # 使用yaml配置文件
# 多进程多卡训练(分布式训练)
python example.py --config_file configs/config.yaml --distributed # 使用yaml配置文件
- 目标支持的backbone有：resnet[18,34,50,101], ,mobilenet_v2等，详见backbone等 ，其他backbone可以自定义添加
- 训练参数可以通过两种方法指定: (1) 通过argparse命令行指定 (2)通过config.yaml配置文件，当存在同名参数时，以配置文件为默认值
| 参数 | 类型 | 参考值 | 说明 |
|---|---|---|---|
| train_data | str, list | - | 训练数据文件，可支持多个文件 |
| test_data | str, list | - | 测试数据文件，可支持多个文件 |
| work_dir | str | work_space | 训练输出工作空间 |
| net_type | str | resnet18 | backbone类型,{resnet,resnest,mobilenet_v2,...} |
| input_size | list | [128,128] | 模型输入大小[W,H] |
| batch_size | int | 32 | batch size |
| lr | float | 0.1 | 初始学习率大小 |
| optim_type | str | SGD | 优化器，{SGD,Adam} |
| loss_type | str | CELoss | 损失函数 |
| scheduler | str | multi-step | 学习率调整策略，{multi-step,cosine} |
| milestones | list | [30,80,100] | 降低学习率的节点，仅仅scheduler=multi-step有效 |
| momentum | float | 0.9 | SGD动量因子 |
| num_epochs | int | 120 | 循环训练的次数 |
| num_warn_up | int | 3 | warn_up的次数 |
| num_workers | int | 12 | DataLoader开启线程数 |
| weight_decay | float | 5e-4 | 权重衰减系数 |
| gpu_id | list | [ 0 ] | 指定训练的GPU卡号，可指定多个 |
| log_freq | in | 20 | 显示LOG信息的频率 |
| finetune | str | model.pth | finetune的模型 |
| use_prune | bool | True | 是否进行模型剪枝 |
| progress | bool | True | 是否显示进度条 |
| distributed | bool | False | 是否使用分布式训练 |
一个简单分类例子如下：
# -*-coding: utf-8 -*-
"""
@Author : panjq
@E-mail : [email protected]
@Date : 2021-07-28 22:09:32
"""
import os
import sys
sys.path.append(os.getcwd())
import argparse
import basetrainer
from torchvision import transforms
from torchvision.datasets import ImageFolder
from basetrainer.engine import trainer
from basetrainer.engine.launch import launch
from basetrainer.criterion.criterion import get_criterion
from basetrainer.metric import accuracy_recorder
from basetrainer.callbacks import log_history, model_checkpoint, losses_recorder, multi_losses_recorder
from basetrainer.scheduler import build_scheduler
from basetrainer.optimizer.build_optimizer import get_optimizer
from basetrainer.utils import log, file_utils, setup_config, torch_tools
from basetrainer.models import build_models
print(basetrainer.__version__)
class ClassificationTrainer(trainer.EngineTrainer):
""" Training Pipeline """
def __init__(self, cfg):
super(ClassificationTrainer, self).__init__(cfg)
torch_tools.set_env_random_seed()
cfg.model_root = os.path.join(cfg.work_dir, "model")
cfg.log_root = os.path.join(cfg.work_dir, "log")
if self.is_main_process:
file_utils.create_dir(cfg.work_dir)
file_utils.create_dir(cfg.model_root)
file_utils.create_dir(cfg.log_root)
file_utils.copy_file_to_dir(cfg.config_file, cfg.work_dir)
setup_config.save_config(cfg, os.path.join(cfg.work_dir, "setup_config.yaml"))
self.logger = log.set_logger(level="debug",
logfile=os.path.join(cfg.log_root, "train.log"),
is_main_process=self.is_main_process)
# build project
self.build(cfg)
self.logger.info("=" * 60)
self.logger.info("work_dir :{}".format(cfg.work_dir))
self.logger.info("config_file :{}".format(cfg.config_file))
self.logger.info("gpu_id :{}".format(cfg.gpu_id))
self.logger.info("main device :{}".format(self.device))
self.logger.info("num_samples(train):{}".format(self.num_samples))
self.logger.info("num_classes :{}".format(cfg.num_classes))
self.logger.info("mean_num :{}".format(self.num_samples / cfg.num_classes))
self.logger.info("=" * 60)
def build_optimizer(self, cfg, **kwargs):
"""build_optimizer"""
self.logger.info("build_optimizer")
self.logger.info("optim_type:{},init_lr:{},weight_decay:{}".format(cfg.optim_type, cfg.lr, cfg.weight_decay))
optimizer = get_optimizer(self.model,
optim_type=cfg.optim_type,
lr=cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
return optimizer
def build_criterion(self, cfg, **kwargs):
"""build_criterion"""
self.logger.info("build_criterion,loss_type:{},num_classes:{}".format(cfg.loss_type, cfg.num_classes))
criterion = get_criterion(cfg.loss_type, cfg.num_classes, device=self.device)
return criterion
def build_train_loader(self, cfg, **kwargs):
"""build_train_loader"""
self.logger.info("build_train_loader,input_size:{}".format(cfg.input_size))
transform = transforms.Compose([
transforms.Resize([int(128 * cfg.input_size[1] / 112), int(128 * cfg.input_size[0] / 112)]),
transforms.RandomHorizontalFlip(),
transforms.RandomCrop([cfg.input_size[1], cfg.input_size[0]]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
dataset = ImageFolder(root=cfg.train_data, transform=transform)
cfg.num_classes = len(dataset.classes)
cfg.classes = dataset.classes
loader = self.build_dataloader(dataset, cfg.batch_size, cfg.num_workers, phase="train",
shuffle=True, pin_memory=False, drop_last=True, distributed=cfg.distributed)
return loader
def build_test_loader(self, cfg, **kwargs):
"""build_test_loader"""
self.logger.info("build_test_loader,input_size:{}".format(cfg.input_size))
transform = transforms.Compose([
transforms.Resize([int(128 * cfg.input_size[1] / 112), int(128 * cfg.input_size[0] / 112)]),
transforms.CenterCrop([cfg.input_size[1], cfg.input_size[0]]),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
dataset = ImageFolder(root=cfg.train_data, transform=transform)
loader = self.build_dataloader(dataset, cfg.batch_size, cfg.num_workers, phase="test",
shuffle=False, pin_memory=False, drop_last=False, distributed=False)
return loader
def build_model(self, cfg, **kwargs):
"""build_model"""
self.logger.info("build_model,net_type:{}".format(cfg.net_type))
model = build_models.get_models(net_type=cfg.net_type, input_size=cfg.input_size,
num_classes=cfg.num_classes, pretrained=True)
if cfg.finetune:
self.logger.info("finetune:{}".format(cfg.finetune))
state_dict = torch_tools.load_state_dict(cfg.finetune)
model.load_state_dict(state_dict)
if cfg.use_prune:
from basetrainer.pruning import nni_pruning
sparsity = 0.2
self.logger.info("use_prune:{},sparsity:{}".format(cfg.use_prune, sparsity))
model = nni_pruning.model_pruning(model,
input_size=[1, 3, cfg.input_size[1], cfg.input_size[0]],
sparsity=sparsity,
reuse=False,
output_prune=os.path.join(cfg.work_dir, "prune"))
model = self.build_model_parallel(model, cfg.gpu_id, distributed=cfg.distributed)
return model
def build_callbacks(self, cfg, **kwargs):
"""定义回调函数"""
self.logger.info("build_callbacks")
# 准确率记录回调函数
acc_record = accuracy_recorder.AccuracyRecorder(target_names=cfg.classes,
indicator="Accuracy")
# loss记录回调函数
loss_record = losses_recorder.LossesRecorder(indicator="loss")
# Tensorboard Log等历史记录回调函数
history = log_history.LogHistory(log_dir=cfg.log_root,
log_freq=cfg.log_freq,
logger=self.logger,
indicators=["loss", "Accuracy"],
is_main_process=self.is_main_process)
# 模型保存回调函数
checkpointer = model_checkpoint.ModelCheckpoint(model=self.model,
optimizer=self.optimizer,
moder_dir=cfg.model_root,
epochs=cfg.num_epochs,
start_save=-1,
indicator="Accuracy",
logger=self.logger)
# 学习率调整策略回调函数
lr_scheduler = build_scheduler.get_scheduler(cfg.scheduler,
optimizer=self.optimizer,
lr_init=cfg.lr,
num_epochs=cfg.num_epochs,
num_steps=self.num_steps,
milestones=cfg.milestones,
num_warn_up=cfg.num_warn_up)
callbacks = [acc_record,
loss_record,
lr_scheduler,
history,
checkpointer]
return callbacks
def run(self, logs: dict = {}):
self.logger.info("start train")
super().run(logs)
def main(cfg):
t = ClassificationTrainer(cfg)
return t.run()
def get_parser():
parser = argparse.ArgumentParser(description="Training Pipeline")
parser.add_argument("-c", "--config_file", help="configs file", default="configs/config.yaml", type=str)
# parser.add_argument("-c", "--config_file", help="configs file", default=None, type=str)
parser.add_argument("--train_data", help="train data", default="./data/dataset/train", type=str)
parser.add_argument("--test_data", help="test data", default="./data/dataset/val", type=str)
parser.add_argument("--work_dir", help="work_dir", default="output", type=str)
parser.add_argument("--input_size", help="input size", nargs="+", default=[224, 224], type=int)
parser.add_argument("--batch_size", help="batch_size", default=32, type=int)
parser.add_argument("--gpu_id", help="specify your GPU ids", nargs="+", default=[0], type=int)
parser.add_argument("--num_workers", help="num_workers", default=0, type=int)
parser.add_argument("--num_epochs", help="total epoch number", default=50, type=int)
parser.add_argument("--scheduler", help=" learning scheduler: multi-step,cosine", default="multi-step", type=str)
parser.add_argument("--milestones", help="epoch stages to decay learning rate", nargs="+",
default=[10, 20, 40], type=int)
parser.add_argument("--num_warn_up", help="num_warn_up", default=3, type=int)
parser.add_argument("--net_type", help="net_type", default="mobilenet_v2", type=str)
parser.add_argument("--finetune", help="finetune model file", default=None, type=str)
parser.add_argument("--loss_type", help="loss_type", default="CELoss", type=str)
parser.add_argument("--optim_type", help="optim_type", default="SGD", type=str)
parser.add_argument("--lr", help="learning rate", default=0.1, type=float)
parser.add_argument("--weight_decay", help="weight_decay", default=0.0005, type=float)
parser.add_argument("--momentum", help="momentum", default=0.9, type=float)
parser.add_argument("--log_freq", help="log_freq", default=10, type=int)
parser.add_argument('--use_prune', action='store_true', help='use prune', default=False)
parser.add_argument('--progress', action='store_true', help='display progress bar', default=True)
parser.add_argument('--distributed', action='store_true', help='use distributed training', default=False)
parser.add_argument('--polyaxon', action='store_true', help='polyaxon', default=False)
return parser
if __name__ == "__main__":
parser = get_parser()
cfg = setup_config.parser_config(parser.parse_args(), cfg_updata=True)
launch(main,
num_gpus_per_machine=len(cfg.gpu_id),
dist_url="tcp://127.0.0.1:28661",
num_machines=1,
machine_rank=0,
distributed=cfg.distributed,
args=(cfg,))
5.可视化
目前训练过程可视化工具是使用Tensorboard，使用方法：
tensorboard --logdir=path/to/log/ |
 |
|---|---|
 |
 |
 |
 |