含float64操作的TensorFlow模型转TFLite遇阻求助

问题概述 #

尝试将谷歌多物种鲸鱼模型转换为TFLite时，尽管模型签名显示所有输入输出均为float32，转换过程仍提示存在FlexMul、FlexRound、FlexSub等float64操作，加载Flex delegate时还出现组件重复注册的冲突错误。

解决方案 #

1. 强制转换模型内部张量为float32 #

通过重定义模型的推理函数，将内部隐式存在的float64张量强制转换为float32，再重新保存模型：

import tensorflow as tf

saved_model_dir = "model"

# 加载原模型并获取目标推理函数
loaded_model = tf.saved_model.load(saved_model_dir)
infer_func = loaded_model.signatures["serving_default"]

# 包装函数，处理张量类型转换
@tf.function(input_signature=infer_func.input_signature)
def converted_func(**kwargs):
    processed_inputs = {}
    # 处理输入：仅转换float64类型为float32，保留其他类型
    for k, v in kwargs.items():
        if v.dtype == tf.float64:
            processed_inputs[k] = tf.cast(v, tf.float32)
        else:
            processed_inputs[k] = v
    outputs = infer_func(**processed_inputs)
    # 处理输出：同样转换float64为float32
    converted_outputs = {}
    for k, v in outputs.items():
        if v.dtype == tf.float64:
            converted_outputs[k] = tf.cast(v, tf.float32)
        else:
            converted_outputs[k] = v
    return converted_outputs

# 保存修改后的模型
tf.saved_model.save(loaded_model, "model_float32", signatures={"serving_default": converted_func})

使用修改后的模型执行TFLite转换：

converter = tf.lite.TFLiteConverter.from_saved_model("model_float32")
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
# 启用默认优化并限定支持类型
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float32]

tflite_model = converter.convert()
with open("converted_model.tflite", "wb") as f:
    f.write(tflite_model)

2. 直接启用TFLite类型强制转换选项 #

无需修改原模型，在转换阶段强制将所有float64操作转换为float32：

import tensorflow as tf

saved_model_dir = "model"

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
# 限定仅支持float32，自动转换float64张量
converter.target_spec.supported_types = [tf.float32]
# 启用资源变量支持，避免转换报错
converter.experimental_enable_resource_variables = True

tflite_model = converter.convert()
with open("converted_model.tflite", "wb") as f:
    f.write(tflite_model)

3. 修复Flex delegate加载冲突错误 #

加载delegate时的注册冲突源于完整TensorFlow环境与TFLite组件的重复注册，解决方法：

• 使用独立Python环境，仅安装tflite_runtime而非完整TensorFlow库
• 编译静态链接版Flex delegate（避免动态库冲突）：

bazelisk build -c opt --config=macos --config=monolithic //tensorflow/lite/delegates/flex:tensorflowlite_flex

使用tflite_runtime加载模型：

from tflite_runtime.interpreter import Interpreter
from tflite_runtime.interpreter import load_delegate

flex_delegate_path = ".../libtensorflowlite_flex.dylib"
flex_delegate = load_delegate(flex_delegate_path)

interpreter = Interpreter(
    model_path="converted_model.tflite",
    experimental_delegates=[flex_delegate]
)
interpreter.allocate_tensors()

4. 清理SavedModel冗余节点 #

部分SavedModel可能包含未被签名使用的float64节点，可通过保留核心签名清理模型：

import tensorflow as tf

saved_model_dir = "model"
loaded_model = tf.saved_model.load(saved_model_dir)

# 仅保留业务需要的签名
signatures_to_keep = ["serving_default", "score", "metadata"]
filtered_signatures = {k: loaded_model.signatures[k] for k in signatures_to_keep}

tf.saved_model.save(loaded_model, "model_cleaned", signatures=filtered_signatures)

使用清理后的模型执行TFLite转换即可。

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