Keras模型转TensorFlow Lite如何规避不支持的运算错误

调整预处理位置规避转换错误 #

报错的根本原因是tf.keras.applications.mobilenet_v2.preprocess_input直接作用于5维时序输入张量时，内部调用的Mul算子触发了TFLite原生算子的广播限制，而单帧4维输入的预处理逻辑是完全被TFLite原生支持的。将预处理和MobileNetV2封装为单帧处理子模型后再套入TimeDistributed层，即可在不修改推理逻辑、不需要重训的前提下解决转换失败问题。

修改后的可正常转换代码：

import tensorflow as tf

IMAGE_SHAPE = (224, 224, 3)

# 封装单帧处理子模型：包含预处理 + MobileNetV2推理
def build_single_frame_model():
    frame_input = tf.keras.Input(shape=IMAGE_SHAPE)
    preprocessed = tf.keras.applications.mobilenet_v2.preprocess_input(frame_input)
    mobilenet = tf.keras.applications.MobileNetV2(
        input_shape=IMAGE_SHAPE,
        include_top=False,
        pooling='avg',
        weights='imagenet'
    )
    frame_output = mobilenet(preprocessed)
    return tf.keras.Model(frame_input, frame_output)

single_frame_model = build_single_frame_model()

# 时序输入处理逻辑
inputs = tf.keras.Input(shape=(5,) + IMAGE_SHAPE)
outputs = tf.keras.layers.TimeDistributed(single_frame_model)(inputs)
model = tf.keras.Model(inputs, outputs)
model.compile()

# 直接转换即可成功，无需启用select ops
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存转换后的模型
with open("mobilenet_temporal.tflite", "wb") as f:
    f.write(tflite_model)

适配说明 #

该方案生成的TFLite模型仅使用原生支持的算子，不需要引入额外的select ops，可直接在基于Mediapipe开发的Android应用中部署使用，推理精度和原结构完全一致。

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