基于强化学习的剪枝调优
基于强化学习的模型剪枝调优包含自动化模型结构解析和基于强化学习的剪枝算法,支持在昇腾910 AI处理器上进行模型搜索和训练,可以用于分类、目标检测、语义分割模型的剪枝压缩,支持以flops(计算量)或params(参数量)或latency(推理时延)为主要目标的模型压缩。在精度无损的情况下,能够有效地提升模型性能。
目前支持对包括但不限于表1列出的模型进行基于强化学习的模型剪枝调优。
类型 |
名称 |
框架 |
---|---|---|
图像分类 |
ResNet50 |
MindSpore、PyTorch |
MobileNetV2 |
MindSpore、PyTorch |
|
图像分割类 |
DeepLabV3 |
MindSpore、PyTorch |
目标检测类 |
FasterRCNN |
MindSpore |
YoloV5 |
MindSpore |
基于强化学习的模型剪枝调优过程
调优过程分为以下两个阶段:
前提条件
- 请执行命令安装graphviz和fvcore包。
- 若待调优模型为MindSpore框架,在剪枝阶段需要根据实际情况配置NPU_VISIBLE_DEVICES和DEVICE_ID两个环境变量。
对MindSpore框架的模型做基于强化学习的模型剪枝调优(backend指定为mindspore)的情况下,对应yml文件需要配置parallel_search和parallel_fully_train为True。
general: backend: mindspore device_category: NPU parallel_search: True parallel_fully_train: True
- 多卡训练时请注意以下几点:
- 环境变量需要增加如下配置:
export HCCL_WHITELIST_DISABLE=1 #关闭HCCL通信白名单 export ORIGIN_RANK_SIZE=8 #服务器总卡数 export BACKEND_TYPE=NPU #训练时的框架类型 export WORLD_SIZE=4 #用于配置vega的进程数量,单机多卡训练下,指定调用卡的数量
- fully_train.pipe_step.type 需要改成HcclTrainStep
pipe_step: type: HcclTrainStep # 单卡使用TrainPipeStep ......
- 数据集配置中的batch_size是指单张卡的batch_size。比如,使用4卡、512batchsize训练模型,yml文件中的batch_size设置为128。
- 环境变量需要增加如下配置:
自定义模型注册
如果需剪枝调优自定义的模型,需要先注册该模型,然后在yml配置使用该模型。目前PyToch框架下DeepLabV3已经注册,ResNet50和MobileNetV2在剪枝调优前需要进行自定义模型注册。MindSpore框架下的ResNet50、MobileNetV2、DeepLabV3均已注册,已注册的模型只需执行1即可。模型注册步骤如下。
- 从 ModelZoo中查找下载模型脚本并放置合适的路径,例如{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/core/networks/pytorch目录下。
- 在模型定义所在脚本文件中使用@Register.register(RegisterType.NETWORK)将网络定义类注册到Vega中,并可能需要对构造函数做一些适配性修改(可以参考DeepLabV3注册代码,例如{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/core/networks/pytorch/deeplab)。
以注册ResNet50模型为例,需要在模型脚本文件(resnet.py)中添加注册代码。
from ascend_automl.tools.register import Register, RegisterType #注册代码 @Register.register(RegisterType.NETWORK) #注册代码 class ResNet50_ModelZoo(Adapter.Module): """ Define Custom Networks. """ def __init__(self, **desc): super(ResNet50_ModelZoo, self).__init__() network_name = desc.get('name', None) network_dir = desc.get('dir', None) class_num = desc.get('class_num', 1001) ......
对{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/module_register.py添加以下模型注册信息。
"ascend_automl.core.networks.pytorch.ResNet50_for_PyTorch.DistributedResnet50.image_classification": { #模型注册脚本相对路径 "resnet": ["ResNet50_ModelZoo"], #"模型注册脚本文件名" : ["模型类名"] },
操作步骤(以ResNet50为例)
- 进入{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/pytorch/prune_amc/classification/resnet50目录,已提供示例文件resnet50_amc.yml(以时延为优化目标的yml文件为同目录下的resnet50_amc_latency.yml),建议拷贝至当前运行目录,根据实际情况配置以下加粗字段(以下代码以resnet50_amc.yml文件为例)。
nas: pipe_step: type: SearchPipeStep model: model_desc: modules: ["backbone"] backbone: type: ResNet_ModelZoo #该模型的定义方法需参照自定义模型注册 version: resnet50 image_size: &image_size 224 #表示模型的输入图像大小 # Warning: pretrained_model_file is possible to construct malicious pickle data which will execute arbitrary code during unpickling pytorch model file. pretrained_model_file: /path/to/pretrained_models/Resnet50.pth # 原模型权重文件,请根据实际路径配置,需确保文件的安全性 dataset: type: Imagenet train: batch_size: 64 #注意,在配置成HcclTrainStep多卡训练时,数据集的batch_size指单张卡的batch_size。例如,使用4卡、512batchsize训练模型,配置中的batch_size设置为128 num_workers: 4 #加载数据时使用的子进程个数 drop_last: False #不足一个batch的数据是否丢弃,取值为“True”或“False” shuffle: True #数据是否随机打乱,取值为“True”或“False” data_path: &dataset_path /path/.../datasets/imagenet/ILSVRC/Data/CLS-LOC # 数据集路径 ... search_algorithm: #请参考剪枝阶段搜索算法配置进行配置 type: PruneAMC objective_keys: 'accuracy' keep_pareto_top_k: 1 plot_train_curve: True plot_graph_save_folder: '/tmp/debug_amc' #临时目录,存储训练过程折线图 # search configs train_episode: 1000 lbound_preserve_ratio: 0.3 reward_type: acc_flops #剪枝任务的目标。acc_flops搜索高精度低flops模型;acc_params搜索高精度低参数量模型;acc_latency搜索高精度低时延模型;acc_abs_latency搜索贴近目标时延的高精度模型 ... trainer: type: Trainer ... optimizer: type: SGD params: lr: 0.001 # 根据实际训练任务设置微调学习率,是正常训练模型的初始学习率的1/10或1/100 momentum: 0.9 cls_fulltrain: pipe_step: type: TrainPipeStep #单卡使用TrainPipeStep,多卡使用HcclTrainStep # type: HcclTrainStep models_folder: '{local_base_path}/output/nas_amc'
- 对图像分类网络模型MobileNetV2进行剪枝压缩时,需要配置的yml文件为{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/pytorch/prune_amc/classification/mobilenet_v2/mobilenetv2_amc.yml(以时延为优化目标的yml文件为mobilenetv2_amc_latency.yml)。
- 对分割网络模型DeepLabV3进行剪枝压缩时,需要配置的yml文件为{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/examples/pytorch/prune_amc/segmentation/deeplabv3/deeplabv3_amc.yml。
- 若模型前向传播不是单个tensor时,请参考模型前向传播不是单个tensor时,如何进行特殊配置进行配置。
- 启动模型剪枝压缩任务。
vega resnet50_amc.yml -d NPU
若剪枝阶段意外中断,请参考剪枝阶段意外中断,如何恢复剪枝搜索恢复剪枝搜索。
剪枝阶段搜索算法配置
剪枝阶段也称为search阶段。强化学习采样剪枝模型,根据采样模型学习每一层的剪枝率,经过多次采样和学习,最终收敛到高精度的剪枝模型。
剪枝算法配置参数含义如表2所示:
参数 |
说明 |
---|---|
keep_pareto_top_k |
选取若干个top精度的剪枝模型作为输出结果。 |
plot_train_curve |
是否画出剪枝算法学习曲线,如果为True,将画出学习曲线并存在output目录下。 |
plot_graph_save_folder |
指定模型graph图像存储目录。 |
train_episode |
剪枝算法采样个数。 |
lbound_preserve_ratio |
0~1的小数,表示剪枝率的下限,算法直接跳过低于该阈值的采样模型。 |
reward_type |
剪枝任务的目标,acc_flops搜索高精度低flops模型;acc_params搜索高精度低参数量模型;acc_latency搜索高精度低时延模型;acc_abs_latency搜索贴近目标时延的高精度模型。 |
target_latency |
当选择时延作为优化目标的时候,需要配置期望的时延,单位ms。 |
channel_round |
模型通道数的最小公倍数。 |
lbound |
每一层通道数保留比例的下限。 |
rbound |
每一层通道数保留比例的上限。 |
seed |
随机种子,可以设置不同随机种子得到不同的采样模型,会产生不同的收敛速度和结果。 |
networkparser |
自动化模型结构解析的配置参数。
|
agent |
强化学习agent参数。
|
specific_bound |
指定某些层的通道保留率的上限和下限,格式{mid: [lbound, rbound]},默认{}。如果lbound等于rbound,即指定当前层剪枝率为固定值。 |
user_skip_module_id |
用户指定某些层不剪。是一个mid list,默认[]。 |