如何将现有TFLite模型中的单个签名保存为仅含单个子图的新模型

你可以通过两种方式实现需求，以下是具体步骤和代码：

方法一：利用TensorFlow高层API快速转换 #

这种方法基于TensorFlow的TFLiteConverter，通过签名对应的concrete function直接生成新模型，操作简单且兼容性好。

import tensorflow as tf
import numpy as np

# 加载原多签名TFLite模型
interpreter = tf.lite.Interpreter(model_path="model.tflite")
# 获取目标签名的运行器
signature_runner = interpreter.get_signature_runner("input_1")
# 提取签名对应的concrete function
concrete_func = signature_runner.concrete_function

# 将concrete function转换为单签名TFLite模型
converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
# 可选：添加优化配置，比如启用默认优化
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
new_tflite_model = converter.convert()

# 保存新模型
with open("single_signature_model.tflite", "wb") as f:
    f.write(new_tflite_model)

验证新模型 #

# 加载新模型并测试
new_interpreter = tf.lite.Interpreter(model_content=new_tflite_model)
new_interpreter.allocate_tensors()

# 获取输入输出张量信息
input_details = new_interpreter.get_input_details()
output_details = new_interpreter.get_output_details()

# 设置测试输入
new_interpreter.set_tensor(input_details[0]['index'], np.array([0,0,0,0,0,0,0,0], dtype=np.float32))
new_interpreter.set_tensor(input_details[1]['index'], np.array([1,2,3,4,5,6,7,8], dtype=np.float32))

# 执行推理
new_interpreter.invoke()

# 获取并打印输出
output = new_interpreter.get_tensor(output_details[0]['index'])
print(output)

方法二：直接操作TFLite模型结构（底层控制） #

如果你需要更精细的模型结构控制，可以直接使用TFLite的Model API提取目标子图并构建新模型：

import tflite

# 读取原模型数据
with open("model.tflite", "rb") as f:
    model_data = f.read()
# 解析原模型
model = tflite.Model.GetRootAsModel(model_data, 0)

# 查找目标签名对应的子图索引
target_signature = "input_1"
target_subgraph_idx = -1
signature_defs = model.SignatureDefs()
for i in range(signature_defs.Length()):
    sig = signature_defs.Get(i)
    if sig.Key().decode() == target_signature:
        target_subgraph_idx = sig.Value().SubgraphIndex()
        break

if target_subgraph_idx == -1:
    raise ValueError(f"未找到签名 {target_signature}")

# 提取目标子图
target_subgraph = model.Subgraphs(target_subgraph_idx)

# 构建新模型
new_model = tflite.Model()
new_model.SetVersion(model.Version())
new_model.SetOperatorCodes(model.OperatorCodes())
# 将目标子图添加到新模型
new_subgraph = new_model.Subgraphs().Add()
new_subgraph.CopyFrom(target_subgraph)
# 清空签名定义（新模型仅含单个子图，无需多签名）
new_model.SetSignatureDefs(tflite.SignatureDefT())

# 序列化并保存新模型
new_model_data = new_model.SerializeToString()
with open("single_subgraph_model.tflite", "wb") as f:
    f.write(new_model_data)

方法对比 #

• 方法一：无需了解TFLite底层结构，代码简洁，自动处理输入输出映射，适合大多数常规场景。
• 方法二：直接操作模型结构，可自定义修改子图细节，适合需要深度定制的场景，但需要熟悉TFLite的模型格式规范。

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