如何让TFLite转换器量化时截断代表性数据集的分布尾部?
问题
我正在将模型量化为int8以在Coral EdgeTPU上运行,使用TFLite转换器的代码如下:
class TensorFlowLiteConverter: def __init__(self, original_network, input_data_generator, tflite_model_folder, tflite_model_file = '', edgetpu_compiler_path = '/bin/edgetpu_compiler' ): """ :param original_network : original keras model. :param input_data_generator : should be a tf.lite.RepresentativeDataset adapted to generate the input for the specified model. :param tflite_model_folder : path to the folder where the tflite model should be saved. :param tflite_model_file : label of the model added to the file name. :param edgetpu_compiler_path: path to the edge tpu compiler, may change depending on the system you are running on. """ self.original_model = original_network self.input_data_generator = input_data_generator self.tflite_model_folder = tflite_model_folder self.tflite_model_file = tflite_model_file self.edgetpu_compiler_path = edgetpu_compiler_path self.tflite_model = None # Used afterwards to compile the model for the tpu def convert_model(self): """ Conversion function from TensorFlow to TensorFlow Lite model. """ converter = tf.lite.TFLiteConverter.from_keras_model(self.original_model) converter.experimental_new_converter = True converter.optimizations = [tf.lite.Optimize.DEFAULT] converter.representative_dataset = self.input_data_generator converter.target_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.inference_input_type = tf.uint8 converter.inference_output_type = tf.uint8 self.tflite_model = converter.convert() self._save_converted_model() return self.tflite_model def _save_converted_model(self): tflite_model_file = self.tflite_model_file + '_int' + time_string + '.tflite' with open(os.path.join(self.tflite_model_folder, tflite_model_file), 'wb') as tflite_out_file: tflite_out_file.write(self.tflite_model) self.tflite_model_file = tflite_model_file print('Tflite model saved as ', os.path.join(self.tflite_model_folder, self.tflite_model_file)) return tflite_model_file def compile_to_edgetpu(self): """ Function to call the EdgeTPU compiler on the saved TFlite integer file. Log files in the model folder give information on the operations mapped to the TPU. """ if not re.search('\\.tflite$', self.tflite_model_file): print('Error, input file must end in .tflite') sys.exit(1) command_arguments = (self.edgetpu_compiler_path, '-o', self.tflite_model_folder, '-m', '13', '-s', os.path.join(self.tflite_model_folder, self.tflite_model_file)) log_file_name = re.sub('\\.tflite$', '.log', os.path.basename(self.tflite_model_file)) assert log_file_name != os.path.basename(self.tflite_model_file) results = subprocess.run(command_arguments, capture_output=True, check=True) with open(os.path.join(self.tflite_model_folder, log_file_name), 'wt') as log_file: log_file.write(results.stderr.decode()) log_file.write(results.stdout.decode()) edgetpu_file = self.tflite_model_file.replace('.tflite', edgetpu_ext, 1) model_path = os.path.join(self.tflite_model_folder , edgetpu_file) print('TPU compiled model saved as ', model_path) return model_path
在使用训练集300个样本构建代表性数据集时遇到问题:样本数量越多,模型内部激活范围越大,量化模型的范围也随之扩大。推测TFLite转换器未截断代表性数据集推理生成的分布尾部,异常值会降低整数模型的表征能力。例如单一层模型中,99.9%输出值在[-1,1]区间,仅一个值为100,希望量化过程忽略该异常值,以更好地表征多数样本。
我尝试使用TFLiteConverterBaseV1的quantized_input_stats属性,以及通过from_saved_model函数设置该属性,但均无法改变范围;而减小代表性数据集规模时,模型输出范围缩小,量化效果更好。请问是否可通过TensorFlow转换器控制该行为?
解决方案
1. 预处理代表性数据集,过滤异常样本
直接修改代表性数据集生成器,在提供样本前过滤掉会触发模型异常激活的样本。比如在数据生成阶段,先对输入样本进行推理,记录模型关键层的激活值,剔除那些导致激活值超出正常范围的样本。这种方法从源头排除异常值,确保用于量化校准的都是正常分布的样本。
2. 使用量化校准的自定义范围截断逻辑
TensorFlow Lite提供**量化感知训练(QAT)**和调试工具来手动约束量化范围:
- 启用量化感知训练,在训练过程中模拟量化行为,让模型的激活分布更稳定,减少异常值的影响。
- 对于预训练模型,使用
tf.lite.experimental.QuantizationDebugger工具分析校准过程中的激活分布,手动设置各层的量化范围(最小值/最大值),强制转换器忽略尾部异常值。示例代码如下:
debugger = tf.lite.experimental.QuantizationDebugger( converter=converter, debug_dataset=input_data_generator, layer_debug_info={} ) # 分析各层激活分布后,手动指定目标层的量化范围 debugger.set_layer_quantization_params( layer_name="your_target_layer", min_val=-1.0, max_val=1.0 ) quantized_model = debugger.convert()
3. 手动传入校准统计数据(TensorFlow 2.10+)
从TensorFlow 2.10版本开始,TFLite转换器支持通过converter.experimental_calibration_stats参数手动传入校准统计数据。你可以自行计算正常样本的激活范围(比如取99.9%分位数作为最大/最小值),然后将这些统计数据传入转换器,替代自动校准的结果。
4. 优化代表性数据集抽样方式
采用分层抽样从训练集中抽取样本,按类别、场景等维度分层,确保样本覆盖主要数据分布,而非简单随机抽取。这样既能保证校准的准确性,又能避免引入异常值,同时不需要过度缩小数据集规模。
注意:quantized_input_stats仅对输入层的量化范围生效,无法控制模型内部各层的激活范围,这也是你之前尝试无效的核心原因。要控制内部层的范围,必须使用上述针对校准过程或量化感知训练的方法。
内容的提问来源于stack exchange,提问作者Kilian Tiziano Le Creurer

