8位量化问题：TFLite回归模型预测与原模型无相关性

我将一个基础线性回归模型量化为Int8格式的TFLite模型后，发现预测输出与原模型结果高度不相关。完整操作步骤如下：

• 训练未量化的线性回归模型（效果正常）
• 基于原模型创建量化感知模型
• 将量化感知模型转换为输入输出均为Int8的TFLite模型
• 将验证数据集转换为Int8格式
• 使用量化模型预测，结果偏差极大且与原模型无相关性

使用的数据集为波士顿房价数据集，完整复现代码如下：

import pandas as pd
import numpy as np
from numpy import loadtxt
# load data
dataset = loadtxt('boston_housing.csv', delimiter=",")
# split into inputs and outputs
dataset_x = dataset[:, :-1]
dataset_y = dataset[:, -1]

dataset_x = np.float32(dataset_x )
dataset_y = np.float32(dataset_y )

from sklearn.model_selection import train_test_split
training_dataset_x, test_dataset_x, training_dataset_y, test_dataset_y = train_test_split(dataset_x, dataset_y, test_size = 0.20)

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
model = Sequential()
model.add(Dense(100, input_dim = 13, activation='relu'))
model.add(Dense(1,))
model.compile(optimizer='rmsprop', loss='mse', metrics=['mae'])
hist = model.fit(training_dataset_x, training_dataset_y, batch_size=32, epochs=500, validation_split=0.2)

#Quantize the model
import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model
q_aware_model = quantize_model(model)

# 'quantize_model' requires a recompile
q_aware_model.compile(optimizer='rmsprop',
                  loss='mse',
                  metrics=['mae'])

q_aware_model.summary()
history = q_aware_model.fit(training_dataset_x, training_dataset_y, batch_size=32, epochs=500, validation_split=0.2)

print(test_dataset_y[1])
print(q_aware_model.predict(test_dataset_x[1].reshape(1, -1)))

#Convert the model to TFLite
import tensorflow as tf

# Create a converter
converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)

# Indicate that you want to perform default optimizations,
# which include quantization
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# Define a generator function that provides your test data's numpy arrays
def representative_data_gen():
    for i in range(10500):
        yield [test_dataset_x[i:i+1]]

# Use the generator function to guide the quantization process
converter.representative_dataset = representative_data_gen

# Ensure that if any ops can't be quantized, the converter throws an error
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# Set the input and output tensors to int8
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

# Convert the model
tflite_model = converter.convert()

# Save the model to disk
open("q_aware_model.tflite", "wb").write(tflite_model)

#Testing the quantized model

# Load the TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="q_aware_model.tflite")
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print(input_details)
print(output_details )

test_x1 = (test_dataset_x-128).astype(np.int8)
print(test_x1)

predictions=[]
for i in range(len(test_x1)):
    test_values = np.expand_dims(test_x1[i].flatten(), axis=0)

    # Set the value for the input tensor
    interpreter.set_tensor(input_details[0]['index'], test_values)

    # Run the inference
    interpreter.invoke()

    output = interpreter.get_tensor(output_details[0]['index'])
    predictions.append(output)

print(predictions) ### prediction values = 50-127, expected values 0-40

问题根源 #

1. 输入数据转换错误：直接用(test_dataset_x-128).astype(np.int8)转换输入不符合TFLite量化模型要求，必须使用模型内置的量化参数（scale和zero_point）做转换，而非固定减128。
2. 代表性数据集越界：representative_data_gen循环10500次，但波士顿数据集测试集仅约101个样本，数组越界导致量化过程使用无效数据，模型失真。
3. 输出未反量化：Int8模型输出是量化后的整数，需要用输出的scale和zero_point转换回浮点值，才能和原模型结果对比。

修正方案 #

1. 修复代表性数据集生成器

def representative_data_gen():
    # 取训练集的有效样本（比如100个），避免越界
    for i in range(min(100, len(training_dataset_x))):
        yield [training_dataset_x[i:i+1]]

2. 正确转换输入数据

# 获取输入量化参数
input_scale = input_details[0]['quantization_parameters']['scales'][0]
input_zero_point = input_details[0]['quantization_parameters']['zero_points'][0]

# 按量化公式转换输入
test_x1 = np.round(test_dataset_x / input_scale + input_zero_point).astype(np.int8)

3. 反量化输出结果

# 获取输出量化参数
output_scale = output_details[0]['quantization_parameters']['scales'][0]
output_zero_point = output_details[0]['quantization_parameters']['zero_points'][0]

predictions=[]
for i in range(len(test_x1)):
    test_values = np.expand_dims(test_x1[i].flatten(), axis=0)
    interpreter.set_tensor(input_details[0]['index'], test_values)
    interpreter.invoke()
    output_int8 = interpreter.get_tensor(output_details[0]['index'])
    # 反量化为浮点值
    output_float = (output_int8 - output_zero_point) * output_scale
    predictions.append(output_float)

4. 额外优化建议

• 量化感知训练时，优先用训练集作为代表性数据集，能提升量化精度。
• 可适当调整量化感知训练的epoch数，避免过拟合或欠拟合。

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