Keras转TFLite模型安卓Kotlin部署推理输出NaN问题求助

问题背景 #

基于Keras搭建二分类卷积神经网络，Python环境下推理结果正常，转换为TFLite格式后在Android端通过Kotlin部署时，推理输出结果全为NaN，同时存在疑问：正常运行的模型输入多为DataType.UINT8，当前模型输入为DataType.FLOAT32，该差异是否为异常诱因。

模型结构与训练代码 #

model = tf.keras.models.Sequential([ 
    tf.keras.layers.Conv2D(16,(3,3), activation = 'relu', input_shape=(size,size,3)),
    tf.keras.layers.MaxPooling2D(2,2),

    tf.keras.layers.Conv2D(32,(3,3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2,2),

    tf.keras.layers.Conv2D(64,(3,3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2,2),

    tf.keras.layers.Conv2D(128,(3,3), activation = 'relu'),
    tf.keras.layers.MaxPooling2D(2,2),

    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation = 'relu'),
    tf.keras.layers.Dense(512, activation = 'relu'),
    tf.keras.layers.Dense(2, activation='softmax',name='predictions')
  ])

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy']) 

Python端推理验证 #

执行推理代码：

print(model.predict(x))

得到正常输出：

[[1. 0.]]

TFLite转换代码 #

# Convert the model.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# Save the model.
with open('model.tflite', 'wb') as f:
  f.write(tflite_model)

Android端部署代码与异常现象 #

Kotlin部署代码如下：

val resizeBitmap: Bitmap = Bitmap.createScaledBitmap(bitmap, 300, 300, true)
val model = com.example.iscanner.ml.Model.newInstance(context)
val theBuffer = TensorImage.fromBitmap(resizeBitmap)
val byteBuffer = theBuffer.buffer
val inputFeature0 = TensorBuffer.createFixedSize(intArrayOf(1, 150, 150, 3), DataType.FLOAT32)
inputFeature0.loadBuffer(byteBuffer)
val outputs = model.process(inputFeature0)
val outputFeature0 = outputs.outputFeature0AsTensorBuffer
val tab = outputFeature0.floatArray
model.close()
Log.d("Tab_test[0]", " "+ tab[0])
Log.d("Tab_test[1]", " "+ tab[1])

运行后日志输出全为NaN：

D/Tab_test[0]:  NaN
D/Tab_test[1]:  NaN

问题根源 #

输出NaN和输入是FLOAT32还是UINT8没有直接关系，故障来自三个明确的输入不匹配问题：

• 输入尺寸不匹配：从Android端定义输入张量的代码可知模型要求输入尺寸为150150，但代码中先把Bitmap缩放到了300300，加载到150*150的输入张量时缓冲区数据长度不匹配，会导致内存读取错位。
• 输入归一化缺失：训练这类图像分类模型时，通常会将0-255范围的像素值归一化到0-1（或对应训练时使用的数值区间），TensorImage.fromBitmap()默认读取的是0-255范围的原始像素值，直接送入FLOAT32模型时数值范围远超训练时的输入分布，卷积计算过程中数值溢出，最终输出NaN。
• 数据类型解析错误：TensorImage.fromBitmap()默认生成UINT8格式的字节缓冲区，直接将UINT8字节流加载到FLOAT32类型的TensorBuffer中，会把字节数据错误解析为浮点数，输入数值完全错乱。

修复方案 #

按以下步骤调整代码即可解决问题：

• 修正Bitmap缩放尺寸，和模型输入要求保持一致，将原来缩放为300300的逻辑改为150150：

val resizeBitmap: Bitmap = Bitmap.createScaledBitmap(bitmap, 150, 150, true)

• 正确处理FLOAT32输入的类型转换与归一化，和Python端训练时的预处理逻辑对齐：

val model = com.example.iscanner.ml.Model.newInstance(context)
// 明确指定输入张量类型为FLOAT32
val inputImage = TensorImage(DataType.FLOAT32)
inputImage.load(resizeBitmap)
// 执行归一化：将0-255像素值缩放到0-1区间，如果训练时用了其他预处理逻辑（比如像素值映射到-1~1），对应调整NormalizeOp参数即可
val imageProcessor = ImageProcessor.Builder()
    .add(NormalizeOp(0f, 255f))
    .build()
val processedInput = imageProcessor.process(inputImage)
// 直接传入处理完成的输入张量推理，不需要手动创建TensorBuffer避免尺寸/类型不匹配
val outputs = model.process(processedInput)
val outputResult = outputs.outputFeature0AsTensorBuffer.floatArray
model.close()
Log.d("Tab_test[0]", " ${outputResult[0]}")
Log.d("Tab_test[1]", " ${outputResult[1]}")

• （可选优化）如果希望使用UINT8类型输入减小模型体积、提升端侧推理速度，可以在TFLite转换阶段开启默认量化：

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

量化后的模型输入为UINT8类型，不需要额外做浮点归一化操作，直接加载Bitmap生成的缓冲区即可完成推理。

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