使用 vai_q_tensorflow2 API 执行量化 - 3.5 简体中文

以下代码展示了如何利用 vai_q_tensorflow2 API 执行训练后量化。您可在此处找到完整示例。

model = tf.keras.models.load_model('float_model.h5')
from tensorflow_model_optimization.quantization.keras import vitis_quantize
quantizer = vitis_quantize.VitisQuantizer(model)
quantized_model = quantizer.quantize_model(calib_dataset=calib_dataset, 
                                           calib_steps=100, 
                                           calib_batch_size=10,
                                           **kwargs) 

calib_dataset

calib_dataset 是校准进程中使用的代表性校准数据集。它可以全部或部分衍生自 eval_dataset、 train_dataset 或其他数据集。

calib_steps

calib_steps 表示校准步骤总数。默认设为 None。如果 calib_dataset 为 tf.data dataset、生成器或 keras.utils.Sequence 实例且 steps 为 None，校准会持续直至数据集完全耗尽为止。阵列输入不支持此实参。

calib_batch_size

calib_batch_size 用于判定校准期间使用的每批次样本数。如果 calib_dataset 为数据集、生成器或 keras.utils.Sequence 实例形式，则批次大小由数据集本身控制。但如果 calib_dataset 为 numpy.array 对象形式，则默认批次大小为 32。

input_shape

list(int)、list(list(int))、tuple(int) 或 dictionary(int)。包含每个输入层的输入形状。如不显式设置，则使用输入层中的默认形状信息。使用多个输入的形状的列表，例如，input_shape=[1, 224, 224, 3] 或 input_shape=[[None, 224, 224, 3], [None, 64, 1]]。所有维度都应包含具体的值，且 batch_size 维度应为 None 或 int。如果模型的输入形状是类似 [None, None, None, 3] 的变量，那么您需指定类似 [None, 224, 224, 3] 的形状，以便生成最终量化模型。

**kwargs

**kwargs 表示用户定义的量化策略配置字典。它支持您改写默认内置量化策略。例如，设置 bias_bit=16 会量化含 16 位量化器的所有偏差。如需了解有关用户定义的配置的更多信息，请参阅 vai_q_tensorflow2 用法 部分。