基于训练脚本的剪枝调优
基于训练脚本的剪枝调优支持在昇腾910 AI处理器上进行模型搜索和训练,覆盖MindSpore,PyTorch框架,能对CV领域分类和检测等常用模型进行自动剪枝调优。
该特性支持以用户的模型训练脚本作为输入,支持自动解析模型结构,对模型进行剪枝以达到性能调优的目的,可以用于分类、检测模型的剪枝压缩,支持以params(参数量)或latency(推理时延)为主要目标的模型压缩。在精度无损的情况下,能够有效地提升模型性能。
目前支持对包括但不限于表1列出的模型进行基于训练脚本的剪枝调优。
类型 |
名称 |
框架 |
---|---|---|
图像分类 |
AlexNet |
PyTorch、MindSpore |
DenseNet121 |
PyTorch |
|
DenseNet161 |
PyTorch |
|
DenseNet169 |
PyTorch |
|
DenseNet201 |
PyTorch |
|
MobileNetV1 |
PyTorch、MindSpore |
|
MobileNetV2 |
PyTorch、MindSpore |
|
ResNet101 |
PyTorch、MindSpore |
|
ResNet152 |
PyTorch |
|
ResNet18 |
PyTorch、MindSpore |
|
ResNet34 |
PyTorch、MindSpore |
|
ResNet50 |
PyTorch、MindSpore |
|
ResNext101 |
PyTorch |
|
ResNext50 |
PyTorch |
|
SqueezeNet1_0 |
PyTorch |
|
SqueezeNet1_1 |
PyTorch |
|
Vgg11 |
PyTorch、MindSpore |
|
Vgg11bn |
PyTorch |
|
Vgg13 |
PyTorch、MindSpore |
|
Vgg13bn |
PyTorch |
|
Vgg16 |
MindSpore |
|
Vgg19 |
PyTorch、MindSpore |
|
wide_resent101_2 |
PyTorch |
|
wide_resent50_2 |
PyTorch |
|
目标检测类 |
cascade_rcnn |
PyTorch |
FasterRCNN |
MindSpore |
|
SSD |
PyTorch、MindSpore |
|
Yolor |
PyTorch |
|
YoloV3 |
MindSpore |
|
YoloV3-tiny |
PyTorch、MindSpore |
|
YoloV4 |
PyTorch |
|
YoloV5 |
PyTorch、MindSpore |
基于训练脚本的剪枝调优过程
调优过程包括以下阶段:
图像分类模型调优(以PyTorch框架的ResNext50为例)
- 进入{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/pytorch/prune/classification/resnext50目录,已提供模型定义文件renext50.py,建议拷贝至当前运行目录修改即可。
如果PyTorch样例工程中没有模型定义文件需要新建一个,在该模型定义文件添加get_model()方法,分以下两种情况:
- 原训练脚本网络是直接从torchvision获取,在模型定义文件中添加get_model()方法,返回该模型定义的实例,如resnext50.py:
from torchvision import models def get_model(): return models.__dict__["resnext50_32x4d"]()
- 若原训练脚本网络是用户自定义的,则将该模型定义代码复制到网络定义文件中,并添加get_model()方法,返回该模型定义的实例,以自定义MobileNetV1网络为例,MobileNetV1.py内容如下:
import … class MobileNetV1(nn.Module): … def get_model(): return MobileNetV1()
- 原训练脚本网络是直接从torchvision获取,在模型定义文件中添加get_model()方法,返回该模型定义的实例,如resnext50.py:
- 进入{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/pytorch/prune/classification/resnext50目录,已提供配置文件resnext50_dag.yml,建议拷贝至当前运行目录修改即可。
- fine_tune阶段主要修改以下加粗字段:
finetune: pipe_step: type: TrainPipeStep model: pretrained_model_file: "" # 模型的预训练权重文件,如无权重文件,可以注释掉。若配置预训练权重文件为pth文件,需确保文件的安全性,否则会有反序列化风险 model_desc: type: Script2Vega path: "" # 模型定义文件路径 image_size: 224
- nas阶段主要修改以下加粗字段:
nas: pipe_step: type: SearchPipeStep search_algorithm: type: RandomSearch objective_keys: ['accuracy', 'latency'] #剪枝搜索的取样指标,这里以准确率和时延作为指标 policy: num_sample: 64 #剪枝搜索的取样个数,这里会尝试搜索出64个剪枝模型,依据objective_keys 评选出最终优化出的模型 search_space: type: PruneDAGSearchSpace model: model_desc_file: "{local_base_path}/output/fine_tune/desc_0.json" pretrained_model_file: "{local_base_path}/output/fine_tune/model_0.pth"
- fully_train阶段无特殊修改字段。
- 此外,在每个阶段中,均提供了dataset、trainer、evaluator字段配置,可以根据需要对数据处理、训练参数、优化器、推理指标做详细配置,以下以Cifar10数据集为例:
- dataset中主要修改以下加粗字段:
dataset: type: Cifar10 #数据集类型 common: data_path: /home/xxx/cifar10 # 数据集路径 batch_size: 128 #一次训练所选取的数据样本数量 train: transforms: - type: Resize size: [256, 256] - type: RandomCrop size: [224, 224] - type: RandomHorizontalFlip - type: ToTensor - type: Normalize mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] val: transforms: - type: Resize size: [224, 224] - type: ToTensor - type: Normalize mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] test: transforms: - type: Resize size: [224, 224] - type: ToTensor - type: Normalize mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225]
- trainer中主要修改以下加粗字段:
trainer: type: Trainer with_train: True model_statistics: True epochs: 1 #训练次数 mixup: False optimizer: type: SGD params: lr: 0.01 #基础学习率 momentum: 0.9 #动量 weight_decay: !!float 1e-4 #权重衰减 lr_scheduler: type: MultiStepLR params: milestones: [30,60,90] gamma: 0.1 #学习率衰减 loss: type: CrossEntropyLoss
- evaluator主要有两种类别:
- 在训练服务器进行评估,evaluator类型为HostEvaluator:
evaluator: type: Evaluator host_evaluator: type: HostEvaluator metric: type: accuracy
- 在推理服务器上进行评估,evaluator类型为DeviceEvaluator:
evaluator: type: Evaluator device_evaluator: type: DeviceEvaluator hardware: "Davinci" remote_host: "http://x.x.x.x:xxxx" #远端推理服务器URL,后四位为端口号,如果在推理服务器中执行“vega-config -q sec”的返回值为“True”,请将“http”更改为“https”
- 在训练服务器进行评估,evaluator类型为HostEvaluator:
- 也可以复用其他阶段的配置,下面以复用fine_tune阶段的dataset字段为例:
dataset: ref: fine_tune.dataset
- dataset中主要修改以下加粗字段:
- fine_tune阶段主要修改以下加粗字段:
- 启动AutoML调优任务。
vega resnext50_dag.yml -d NPU
目标检测模型调优(以MindSpore框架的YoloV5为例)
- 准备训练脚本。基于用户准备的训练脚本构造模型,配置模型定义文件,在模型定义文件中添加get_model()方法,返回该模型定义的实例。
- 配置Vega的yml文件。AutoML工具在{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/路径下,./mindspore/prune/detection和./pytorch/prune/detection目录中提供了针对MindSpore和PyTorch模型的配置文件样例,建议拷贝至当前运行目录,根据实际情况修改即可运行。
- 注册接口。关于接口的更多配置信息,请参考{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/docs/trainer/use_original_trainer_v1.1.md。
以yolov5_ms_prune.yml为例,请参考以下加粗字段进行配置修改。
register: pkg_path:["/home/automl/models/yolov5_prune/"] #源码文件的路径,请根据实际情况配置 modules: - module: "src.yolo" #模块导入,请根据实际路径配置 script_network: ["get_model","get_eval_model"] #注册的函数,也支持配置一个类,比如Resnet50,但接口须符合要求 - module: "train" #模块导入 ori_train_func: ["run_train"] #注册的训练函数 - module: "eval" #模块导入 ori_eval_func: ["run_eval"] # 注册的评估函数
- fine_tune阶段主要修改以下加粗字体信息:
finetune: pipe_step: type: TrainPipeStep model: pretrained_model_file: /home/cache/yolov5/pre_train/yolov5_baseline.ckpt #必选,训练好的模型权重文件,请根据实际路径配置 model_desc: type: Script2Vega ori_network_config: type: ScriptModelGen common: multiple_inputs: #必选,模型需要的输入类型和shape - shape: [ 1, 12, 320, 320 ] dtype: fp32 train: network: #已注册的获取模型实例的接口 type: get_model evaluate: network: #已注册的获取模型实例的接口 type: get_eval_model trainer: type: OriTrainer use_dag_forward: False with_train: False #若需要fine_tune,则配置为True ori_trainer: type: run_train #run_train已注册的训练接口 config: #模型训练需要配置的参数 max_epoch: 5 data_dir: /home/cache/datasets/coco-all/data/coco/ per_batch_size: 32 lr: 0.0003 evaluator: type: Evaluator host_evaluator: type: OriHostEvaluator ori_eval: type: run_eval #run_eval已注册精度评估接口 config: #精度评估需要配置的参数 data_dir: /home/cache/datasets/coco-all/data/coco/ per_batch_size: 32 device_evaluator: type: DeviceEvaluator custom: CustomEvaluator save_intermediate_file: True #保存导出的air模型 hardware: "Davinci" remote_host: "http://x.x.x.x:x" muti_input: True
- nas阶段主要修改以下加粗字体信息:
nas: pipe_step: type: SearchPipeStep search_algorithm: type: RandomSearch objective_keys: ['mAP', 'latency'] #必选,配置剪枝搜索的取样指标,当前示例以准确率和时延作为指标 policy: num_sample: 64 #配置剪枝搜索的取样个数,如示例中会尝试搜索出64个剪枝模型,依据objective_keys评选出最终优化后模型 search_space: type: SCOPDAGSearchSpace hyperparameters: - key: prune_d_rate type: INT range: [ 75, 95 ] #通道数保留率
- fully_train阶段主要修改以下加粗字体信息:
fully_train: pipe_step: type: TrainPipeStep models_folder: "{local_base_path}/output/nas/" trainer: #根据实际训练情况配置超参 ref: finetune.trainer with_train: True ori_trainer: type: run_train config: max_epoch: 60 data_dir: /home/cache/datasets/coco-all/data/coco/ per_batch_size: 32 lr: 0.0003 warmup_epochs: 0
- 注册接口。关于接口的更多配置信息,请参考{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/docs/trainer/use_original_trainer_v1.1.md。
- 配置环境变量。在运行的训练环境中将ascend_automl/的父目录加入PYTHONPATH中:
export PYTHONPATH={CANN包安装路径}/ascend-toolkit/latest/tools:$PYTHONPATH
- 启动AutoML调优任务。
vega yolov5_ms_prune.yml -d NPU
获取调优后模型
- onnx格式模型输出(仅支持PyTorch框架)。
基于训练脚本的模型自动调优支持输出调优后模型的onnx格式,操作参考如下(以PyTorch框架的resnext50样例为例):
- 支持用户在训练脚本中加载调优后模型(支持PyTorch框架和MindSpore框架)。
调优后的网络描述以json文件格式保存,用户可在自己的训练工程中,依据此文件加载优化后的网络及权重。AutoML工具提供load函数,此函数以ori_net(原网络)、desc_path(调优后网络描述文件)、weight_path(调优后网络权重文件)作为输入。
操作参考如下:
- 生成load函数,有以下两种方式:
- 在yml文件的trainer.callbacks中进行配置。
- 执行以下命令。
python3 -m ascend_automl.tools.create_network_loader output_dir backend
- output_dir为生成的代码文件所在目录,须为存在目录且路径下没有名为“network_loader_scripts”文件夹。
- backend为框架,须为pytorch或mindspore,运行结束后,会在output_dir /network_loader_scripts生成network_loader_user.py和utils_user.py。
- 使用load函数加载调优后模型。
用户需在训练工程中添加代码,参考示例如下,运行前需将network_loader_scripts文件夹所在路径配置在PYTHONPATH中:
from network_loader_scripts.network_loader_user import load # 从network_loader_user导入load函数 from network import Net # 假设此处为用户原网络的导入 net = Net() # 假设net为用户原网络 desc = '{task_id}/output/fully_train/desc_2.json' # 调优后网络描述文件,请根据实际路径替换 weight = '{task_id}/output/fully_train/model_2.pth' # 调优后网络权重文件,请根据实际路径替换 pruned_net = load(net, desc, weight) # 获取优化后网络(加载权重) # pruned_net = load(net, desc) # 获取优化后网络(不加载权重)
pruned_net为加载到的调优后模型。
对模型进行剪枝调优后,若调优结果仍达不到预期,可通过量化过程进一步调优,操作步骤可参考基于强化学习的量化调优。
- 生成load函数,有以下两种方式: