目录
Pytorch基础训练库Pytorch-Base-Trainer(PBT)(支持分布式训练)
1.Introduction
2.Install
3.训练框架
(1)训练引擎(Engine)
(2)回调函数(Callback)
4.使用方法
5.Example: 构建自己的分类Pipeline
6.可视化
7.其他
尊重原创,转载请注明出处:https://panjinquan.blog.csdn.net/article/details/122662902
GitHub - PanJinquan/Pytorch-Base-Trainer: Pytorch分布式训练框架Pytorch分布式训练框架. Contribute to PanJinquan/Pytorch-Base-Trainer development by creating an account on GitHub.https://github.com/PanJinquan/Pytorch-Base-Trainer
考虑到深度学习训练过程都有一套约定成俗的流程,鄙人借鉴Keras开发了一套基础训练库: Pytorch-Base-Trainer(PBT); 这是一个基于Pytorch开发的基础训练库,支持以下特征:
诚然,诸多大公司已经开源基础库,如MMClassification,MMDetection等库; 但碍于这些开源库安装麻烦,依赖库多,版本差异大等问题;鄙人开发了一套比较基础的训练Pipeline:Pytorch-Base-Trainer(PBT), 基于PBT可以快速搭建自己的训练工程; 目前,基于PBT完成了通用分类库(PBTClassification),通用检测库(PBTDetection),通用语义分割库( PBTSegmentation)以及,通用姿态检测库(PBTPose)
通用库 | 类型 | 说明 |
---|---|---|
PBTClassification | 通用分类库 | 集成常用的分类模型,支持多种数据格式,样本重采样 |
PBTDetection | 通用检测库 | 集成常用的检测类模型,如RFB,SSD和YOLOX |
PBTSegmentation | 通用语义分割库 | 集成常用的语义分割模型,如DeepLab,UNet等 |
PBTPose | 通用姿态检测库 | 集成常用的人体姿态估计模型,如UDP,Simple-base-line |
基于PBT框架训练的模型,已经形成了一套完整的Android端上部署流程,支持CPU和GPU
人体姿态估计2DPose | 人脸+人体检测 | 人像抠图 |
---|---|---|
![]() |
||
CPU/GPU:70/50ms | CPU/GPU:30/20ms | CPU/GPU:150/30ms |
PS:受商业保护,目前,仅开源Pytorch-Base-Trainer(PBT),基于PBT的分类,检测和分割以及姿态估计训练库,暂不开源。
git clone https://github.com/PanJinquan/Pytorch-Base-Trainer
cd Pytorch-Base-Trainer
bash setup.sh #pip install dist/basetrainer-*.*.*.tar.gz
pip install basetrainer
# Linux or macOS
python3 -m pip install --upgrade nni
# Windows
python -m pip install --upgrade nni
PBT基础训练库定义了一个基类(Base),所有训练引擎(Engine)以及回调函数(Callback)都会继承基类。
Engine
类实现了训练/测试的迭代方法(如on_batch_begin,on_batch_end),其迭代过程参考如下,用户可以根据自己的需要自定义迭代过程:
self.on_train_begin()
for epoch in range(num_epochs):
self.set_model() # 设置模型
# 开始训练
self.on_epoch_begin() # 开始每个epoch调用
for inputs in self.train_dataset:
self.on_batch_begin() # 每次迭代开始时回调
self.run_step() # 每次迭代返回outputs, losses
self.on_train_summary() # 每次迭代,训练结束时回调
self.on_batch_end() # 每次迭代结束时回调
# 开始测试
self.on_test_begin()
for inputs in self.test_dataset:
self.run_step() # 每次迭代返回outputs, losses
self.on_test_summary() # 每次迭代,测试结束时回调
self.on_test_end() # 结束测试
# 结束当前epoch
self.on_epoch_end()
self.on_train_end()
EngineTrainer
类继承Engine
类,用户需要继承该类,并实现相关接口:
接口 | 说明 |
---|---|
build_train_loader | 定义训练数据 |
build_test_loader | 定义测试数据 |
build_model | 定义模型 |
build_optimizer | 定义优化器 |
build_criterion | 定义损失函数 |
build_callbacks | 定义回调函数 |
另外,EngineTrainer
类还是实现了两个重要的类方法(build_dataloader和build_model_parallel),用于构建分布式训练
类方法 | 说明 |
---|---|
build_dataloader | 用于构建加载方式,参数distributed设置是否使用分布式加载数据 |
build_model_parallel | 用于构建模型,参数distributed设置是否使用分布式训练模型 |
每个回调函数都需要继承(Callback),用户在回调函数中,可实现对迭代方法输入/输出的处理,例如:
回调函数 | 说明 |
---|---|
LogHistory | Log历史记录回调函数,可使用Tensorboard可视化 |
ModelCheckpoint | 保存模型回调函数,可选择最优模型保存 |
LossesRecorder | 单个Loss历史记录回调函数,可计算每个epoch的平均值 |
MultiLossesRecorder | 用于多任务Loss的历史记录回调函数 |
AccuracyRecorder | 用于计算分类Accuracy回调函数 |
get_scheduler | 各种学习率调整策略(MultiStepLR,CosineAnnealingLR,ExponentialLR)的回调函数 |
basetrainer
使用方法可以参考example.py,构建自己的训练器,可通过如下步骤实现:
ClassificationTrainer
,继承trainer.EngineTrainer
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,))
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配置文件
参数 | 类型 | 参考值 | 说明 |
---|---|---|---|
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,))
目前训练过程可视化工具是使用Tensorboard,使用方法:
tensorboard --logdir=path/to/log/
![]() |
![]() |
---|---|
![]() |
![]() |
![]() |
![]() |
作者 | PKing |
---|---|
联系方式 | [email protected] |