基于训练脚本的剪枝调优支持在Atlas 训练系列产品上进行模型搜索和训练，覆盖MindSpore，PyTorch框架，能对CV领域分类和检测等常用模型进行自动剪枝调优。

该特性支持以用户的模型训练脚本作为输入，支持自动解析模型结构，对模型进行剪枝以达到性能调优的目的，可以用于分类、检测模型的剪枝压缩，支持以params（参数量）或latency（推理时延）为主要目标的模型压缩。在精度无损的情况下，能够有效地提升模型性能。

目前支持对包括但不限于表1列出的模型进行基于训练脚本的剪枝调优。

表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

基于训练脚本的剪枝调优过程

调优过程包括以下阶段：

微调阶段（fine_tune）
剪枝阶段（nas）
1. 采样剪枝模型。
2. 微调训练剪枝模型，评估模型精度。
3. 重复1~2，经过N次采样和学习，剪枝结束，选取若干个top精度的剪枝模型作为输出结果，进行下一步训练。
剪枝后训练阶段（fully_train）
训练剪枝后的模型。

图像分类模型调优（以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()
```
如果选用MindSpore框架的样例，需进入对应的MindSpore目录，从ModelZoo下载对应的模型定义文件，添加get_model()方法，以新建resnet18.py，数据集为Cifar10为例：
```
import …

class resnet18(class_num=10):
…
def get_model():
    return resnet18(10)
```

进入{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”

也可以复用其他阶段的配置，下面以复用fine_tune阶段的dataset字段为例：
```
dataset:
  ref: fine_tune.dataset
```

基于MindSpore框架的训练脚本现在仅支持在Atlas 训练系列产品上做推理，即MindSpore框架的yml文件中的evaluator类型只支持HostEvaluator。
其他参数配置可参考Vega官网。
如果需要用到非cifar10数据集，字段模板可参考Vega官网。

启动AutoML调优任务。
```
vega resnext50_dag.yml -d NPU
```

目标检测模型调优（以MindSpore框架的YoloV5为例）

准备训练脚本。基于用户准备的训练脚本构造模型，配置模型定义文件，在模型定义文件中添加get_model()方法，返回该模型定义的实例。
以yolov5为例，在模型定义文件中参考以下加粗字体信息进行修改。
```
def get_model():
    return YOLOV5(is_training=True)
    
def get_eval_model():
    return YOLOV5(is_training=False)
    ...
```

配置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

配置环境变量。
在运行的训练环境中将ascend_automl/的父目录加入PYTHONPATH中：
```
export PYTHONPATH={CANN包安装路径}/ascend-toolkit/latest/tools:$PYTHONPATH
```
启动AutoML调优任务。
```
vega yolov5_ms_prune.yml -d NPU
```

获取调优后模型

onnx格式模型输出（仅支持PyTorch框架）。
基于训练脚本的模型自动调优支持输出调优后模型的onnx格式，操作参考如下（以PyTorch框架的resnext50样例为例）：
1. 在resnext50_dag.yml文件中nas阶段的trainer字段中配置callbacks字段：
```
callbacks:OnnxModelExportCallback
```
2. 启动AutoML调优任务。
  任务结束后会在{task_id}/output/nas/路径下输出model_{worker_id}.onnx文件。
  
  若用户需将onnx文件转换成PyTorch脚本定义的模型，具体步骤可参考onnx-pytorch。
支持用户在训练脚本中加载调优后模型（支持PyTorch框架和MindSpore框架）。
调优后的网络描述以json文件格式保存，用户可在自己的训练工程中，依据此文件加载优化后的网络及权重。AutoML工具提供load函数，此函数以ori_net（原网络）、desc_path（调优后网络描述文件）、weight_path（调优后网络权重文件）作为输入。

操作参考如下：
1. 生成load函数，有以下两种方式：
  - 在yml文件的trainer.callbacks中进行配置。
    1. 在yml文件中nas阶段的trainer字段中配置callbacks字段：
```
callbacks:NetworkLoaderCreateCallback
```
    2. 启动AutoML调优任务。
      任务结束后会在{task_id}/output/nas/network_loader_scripts路径下生成network_loader_user.py（包含load函数）和utils_user.py文件。
  - 执行以下命令。
```
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。
2. 使用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为加载到的调优后模型。
对模型进行剪枝调优后，若调优结果仍达不到预期，可通过量化过程进一步调优，操作步骤可参考基于强化学习的量化调优。