搜索流程

本节介绍自动混合精度搜索场景的接口调用流程和调用示例。

接口调用流程

接口调用流程如图1所示，蓝色部分为用户实现，灰色部分为用户调用AMCT提供的API实现。

用户准备好PyTorch的原始模型、自动混合精度配置文件和推理评估器（Evaluator），调用auto_mixed_precision_search，根据压缩率、量化位宽、量化敏感度以及计算复杂度信息，执行自动混合精度搜索，得到混合精度配置文件与可用于量化感知训练的简易配置文件。

其中Evaluator模块需要用户自定义，用来执行模型的推理，获取量化因子，dump数据（每一层的输入数据）等信息。

图1 调用示例

调用示例

本示例演示了使用AMCT进行自动混合精度搜索的流程，该过程需要用户实现一个模型推理和校准的评估器。

  • 如下示例标有“由用户补充处理”的步骤，需要用户根据自己的模型和数据集进行补充处理，示例中仅为示例代码。
  • 如下示例调用AMCT的部分，函数入参请根据实际情况进行调整。
  1. 导入AMCT包，并通过AMCT（PyTorch）>AMCT(PyTorch)中的环境变量设置日志级别。

    1
2
    		 
    import amct_pytorch as amct
from amct_pytorch.common.auto_calibration import AutoCalibrationEvaluatorBase

  2. （由用户补充处理）实现一个模型的评估器，使用原始待量化的模型和测试集，实现回调函数calibration()、evaluate()。其中：

    • calibration()完成校准的推理。
    • evaluate()完成模型的前向推理过程。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
    		 
    class ModelEvaluator(AutoCalibrationEvaluatorBase):
    # The evaluator for model
    def __init__(self, *args, **kwargs):
        # 做成员变量初始化 
        pass

    def calibration(self, model, batch_num):
        """ 对模型做量化校准前向推理
        model：torch.nn.Module类型, 对模型model做前向推理
        batch_num: int类型，前向推理的batch数目，与量化配置的batch_num一致
        """
        pass

    def evaluate(self, model, iterations):
        """ 对一张图做量化校准前向推理
        model：torch.nn.Module类型, 对模型model做前向推理
        iterations: int类型，前向推理的batch数目
        """
        pass

  3. 调用AMCT工具，进行自动混合精度搜索。

    1. 初始化evaluator，可以使用2构造的ModelEvaluator，也可以使用工具提供的ModelEvaluator，请参见AMCT安装目录/amct_pytorch/utils/evaluator.py。 
      1
2
3
4
5
6
      		 
      evaluator = ModelEvaluator()
或者
evaluator = amct.ModelEvaluator(
    data_dir="./data/input_bin/", 
    input_shape="input:32,3,224,224", 
    data_types="float32")
    1. 进行自动混合精度搜索。 
      1
2
3
4
5
      		 
      model = build_model()
input_data = tuple([torch.randn(input_shape)])
amc_config = './amc.cfg'
save_dir = './results/model_name'
amct.auto_mixed_precision_search(model=model, input_data=input_data, config=amc_config, save_dir=save_dir, evaluator=evaluator, sensitivity='MseSimilarity')

父主题： 自动混合精度搜索