MNIST准确率99%的CNN转TFLite后输出固定及_absl_警告问题

问题根源 #

你看到的_jit_compiled_convolution_op未追踪警告是核心原因：TensorFlow训练CNN时会自动给卷积层启用即时编译(JIT)优化，但转换TFLite时，这些被JIT编译的层没有被完整序列化，导致TFLite模型里的卷积层计算逻辑或参数异常，最终每次推理都输出同一个结果。移除卷积和池化层后警告消失，也直接印证了问题出在JIT处理的卷积层上。

解决办法 #

1. 禁用JIT后重新转换 #

先关闭TensorFlow的JIT编译，确保所有层都能被正确追踪序列化：

# 关闭JIT编译
tf.config.experimental.set_jit(False)
# 重新加载训练好的Keras模型
model = tf.keras.models.load_model("你的原模型文件.h5")
# 执行TFLite转换
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open("mnist_tflite.tflite", "wb") as f:
    f.write(tflite_model)

2. 开启转换兼容选项 #

如果禁用JIT没用，试试开启TFLite转换的资源变量兼容选项，强制追踪所有层：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 启用资源变量兼容，确保所有函数被完整序列化
converter.experimental_enable_resource_variables = True
tflite_model = converter.convert()
with open("mnist_tflite.tflite", "wb") as f:
    f.write(tflite_model)

3. 训练阶段提前禁用JIT（可选） #

如果后续还要训练新模型，直接在训练前关闭JIT，从根源避免转换问题：

tf.config.experimental.set_jit(False)
# 构建并训练模型
model=tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(64,(3,3),input_shape=(28,28,1),activation=tf.nn.relu),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dropout(0.25),
    tf.keras.layers.Dense(128,activation=tf.nn.relu),
    tf.keras.layers.Dense(10,activation=tf.nn.softmax),
])
# 后续训练代码...

验证转换效果 #

转换完成后，拿几个MNIST测试样本分别输入原Keras模型和TFLite模型，对比输出结果，确认TFLite模型的输出不再固定，且和原模型结果一致即可。

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