TensorFlow转TFLite参数未全量化为int8的原因求助

将TensorFlow模型转换为.tflite格式时，发现生成的.tflite文件中并非所有参数都被量化为int8：仅Dense层参数为int8，而Conv2D层参数仍为float32。但TensorFlow官方文档明确说明参数会被量化为int8，预期所有参数完成量化，而激活部分不量化。

环境信息 #

• OS: Ubuntu 22.04
• GPU: RX 3090
• TensorFlow: 2.4

代码示例 #

import tensorflow as tf
import numpy as np
import pathlib
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

train_images = train_images.astype(np.float32) / 255.0
test_images = test_images.astype(np.float32) / 255.0


model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Flatten(), 
tf.keras.layers.Dense(10) 
])
model.summary()

model.compile(optimizer='adam',
          loss=tf.keras.losses.SparseCategoricalCrossentropy(
              from_logits=True),
          metrics=['accuracy'])
model.fit(
train_images,
train_labels,
epochs=1,
validation_data=(test_images, test_labels)
)

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAUL]

tflite_model_quant = converter.convert()

'''
The code for inspect the tensors:
'''

interpreter = tflite.Interpreter(model_content=tflite_model_quant)
interpreter.allocate_tensors()
tensor_details = interpreter.get_tensor_details()
for tensor in tensor_details:
    tensor_name = tensor['name']
    tensor_shape = tensor['shape']
    tensor_type = tensor['dtype']
    tensor_params = np.prod(tensor_shape)
    print(f"Tensor Name: {tensor_name}")
    print(f"Tensor Shape: {tensor_shape}")
    print(f"Number of Parameters: {tensor_params}")
    print(f"Tensor Type:{tensor_type}")
    print(f"Tensor Parameter:\n{interpreter.get_tensor(tensor['index'])}")

1. 修复代码拼写错误 #

代码中存在笔误：tf.lite.Optimize.DEFAUL应为tf.lite.Optimize.DEFAULT，少了字母t会导致优化配置不生效，这是参数未正确量化的直接原因之一。

2. 全整数量化（强制所有权重为int8） #

默认优化（DEFAULT）属于混合量化，会在资源允许时保留部分浮点参数。若要强制所有权重量化为int8，需使用全整数量化，步骤如下：

• 提供校准数据集，让转换器统计权重分布完成量化校准
• 指定支持int8操作集
• 设置输入输出 dtype 为int8

修改后的转换代码：

# 定义校准数据集生成器
def representative_data_gen():
    for input_value in tf.data.Dataset.from_tensor_slices(train_images).batch(1).take(100):
        yield [input_value]

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 指定支持int8内置操作
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# 设置输入输出为int8
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
# 传入校准数据
converter.representative_dataset = representative_data_gen

tflite_model_quant = converter.convert()

3. 权重仅量化（满足权重int8、激活浮点需求） #

若明确只需要权重量化为int8，激活保持float32，可使用以下配置：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model_weight_quant = converter.convert()

此模式下TensorFlow Lite会自动量化权重为int8，激活保留浮点，但需注意TensorFlow 2.4对该特性的支持有局限，建议升级版本。

4. 版本兼容性建议 #

TensorFlow 2.4的量化功能存在局限性，部分层的量化支持不完善。建议升级到TensorFlow 2.8及以上稳定版本，新版本对Conv2D等层的量化支持更稳定，能更好匹配官方文档描述。

内容的提问来源于stack exchange，提问作者LNCC99