下载
中文
注册

auto_mixed_precision_search

功能说明

根据原始模型和自动混合精度搜索简易配置文件,自动搜索模型的最优混合量化位宽配置,生成量化感知训练简易配置文件。

函数原型

auto_mixed_precision_search(model_file, outputs, config, save_dir, evaluator, sensitivity='MseSimilarity')

参数说明

参数名

输入/返回值

含义

使用限制

model_file

输入

TensorFlow pb格式的推理模型,对该模型搜索混合精度。

数据类型:tf.Graph

outputs

输入

模型输出节点的名称。

数据类型:列表,每个元素类型为string。

config

输入

自动混合精度搜索简易配置文件。

基于basic_info.proto文件中的AutoMixedPrecisionConfig生成的简易配置文件,basic_info.proto件所在路径为:AMCT安装目录/amct_tensorflow/proto/basic_info.proto。

basic_info.proto文件参数解释以及生成的自动混合精度搜索配置文件样例请参见自动(混合精度或通道稀疏)搜索简易配置文件说明

数据类型:string

save_dir

输入

用于保存生成的QAT简易配置文件的路径。

该路径需要包含模型名前缀,例如./auto_mixed_precision/*model

数据类型:string

evaluator

输入

进行校准和评估精度的python实例。

数据类型:AutoCalibrationEvaluatorBase派生类。

使用约束:需实现evaluate方法和calibration方法。

sensitivity

输入

评价每一层量化层对于量化敏感度的指标,默认是MSE(Mean Square Error,均方误差)。

数据类型:string或python instance

默认值:MseSimilarity

返回值说明

无。

函数输出

  • *_qat_mixed_precision.json:混合精度配置文件,描述了模型的每一量化层应该使用哪一种精度,*代表模型名前缀。
  • *_qat_mixed_precision.cfg:用于量化感知训练的混合精度配置文件。
  • 中间件(可选):包括层量化敏感度、计算量文件、量化因子记录文件、量化配置文件以及量化层的输入数据文件。

    上述文件在搜索完成后会被删除,只有set_logging_level中的save_level参数设置为debug模式才会被保留。

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
import amct_tensorflow as amct
# 构建待自动混合精度搜索的模型文件
model_file = './model/user_model.pb'
outputs = ['user_model_outputs0', 'user_model_outputs1']
# 自动混合精度搜索简易配置文件
cfg_file = './configs/auto_mixed_precision.cfg'
# 保存结果路径
save_dir = './results/auto_mixed_precision'
# 建立evaluator,实现evaluate和calibration方法
evaluator = amct.GraphEvaluator(
    data_dir="./data/input_bin", 
    input_shape="input:32,3,224,224", 
    data_types="float32")

amct.auto_mixed_precision_search(
    model_file='./model/user_model.pb', 
    outputs=['user_model_outputs0', 'user_model_outputs1'], 
    config='./configs/auto_mixed_precision.cfg', 
    save_dir='./results/auto_mixed_precision', 
    evaluator=evaluator, 
    sensitivity='MseSimilarity')