基于强化学习的剪枝调优

基于强化学习的模型剪枝调优过程

调优过程分为以下两个阶段：

• 剪枝阶段（nas）
  1. 强化学习代理器采样剪枝模型。
  2. 微调训练剪枝模型，评估模型精度。
  3. 根据评估精度计算反馈值，强化学习更新代理器。
  4. 重复1~3，经过N次采样和学习，剪枝结束，选取若干个top精度的剪枝模型作为输出结果，进行下一步训练。
• 剪枝后训练阶段（fully_train）

  训练剪枝后的模型。

前提条件

• 请执行以下命令安装graphviz和fvcore包。

  pip3 install graphviz --user
  pip3 install fvcore --user

• 若待调优模型为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框架下的ResNet50和MobileNetV2在剪枝调优前需要进行自定义模型注册，DeepLabV3不需要。目前MindSpore框架下的ResNet50、MobileNetV2、DeepLabV3均已注册，只需执行1即可。模型注册步骤如下。

1. 从 ModelZoo中查找下载模型脚本并放置合适的路径，例如{CANN包安装路径}/ascend-toolkit/latest/tools/ascend_automl/core/networks/pytorch目录下。
2. 在模型定义所在脚本文件中使用@Register.register(RegisterType.NETWORK)将网络定义类注册到Vega中，并可能需要对构造函数做一些适配性修改。
   以注册ResNet50模型为例，需要在模型脚本文件（resnet.py）中添加注册代码。

   from ascend_automl.tools.register import Register, RegisterType  #注册代码
   
   @Register.register(RegisterType.NETWORK)   #注册代码
   class ResNet50(Adapter.Module):
       """
       Define Custom Networks.
       """
       def __init__(self, **desc):
           super(ResNet50, 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_1.4_code.DistributedResnet50.image_classification": {  #模型注册脚本相对路径
       "resnet": ["ResNet50"],   #"模型注册脚本文件名" : ["模型类名"]
   },

操作步骤（以ResNet50为例）

1. 进入{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时，如何进行特殊配置进行配置。

2. 启动模型剪枝压缩任务。

   vega resnet50_amc.yml -d NPU

   若剪枝阶段意外中断，请参考剪枝阶段意外中断，如何恢复剪枝搜索恢复剪枝搜索。

剪枝阶段搜索算法配置

剪枝阶段也称为search阶段。强化学习采样剪枝模型，根据采样模型学习每一层的剪枝率，经过多次采样和学习，最终收敛到高精度的剪枝模型。

剪枝算法配置参数含义如表2所示：