多签名场景下TFLite模型输入输出量化的配置方法咨询

方案1：使用TensorFlow 2.10及以上版本原生支持的张量级类型配置 #

从TF 2.10开始，TFLiteConverter新增了inference_input_type_overrides和inference_output_type_overrides两个参数，直接支持按张量名称给不同输入/输出指定量化类型，天然适配多签名、多输入输出场景：

import tensorflow as tf

# 第一步：定义带多签名的模型，这里以Keras模型导出多签名为例
class MultiSigModel(tf.keras.Model):
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 224,224,3], dtype=tf.float32, name="img_input")])
    def classify(self, img_input):
        # 分类逻辑实现
        return self(img_input)
    
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 224,224,3], dtype=tf.float32, name="det_input")])
    def detect(self, det_input):
        # 检测逻辑实现
        return self.backbone(det_input), self.head(det_input)

model = MultiSigModel()
# 导出为SavedModel时保留两个签名
tf.saved_model.save(model, "./saved_model", signatures={"classify": model.classify, "detect": model.detect})

# 第二步：转换时按张量名称指定输入输出量化类型
converter = tf.lite.TFLiteConverter.from_saved_model("./saved_model", signature_keys=["classify", "detect"])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

# 关键配置：按张量名称指定输入类型，两个签名的输入分别叫img_input、det_input
converter.inference_input_type_overrides = {
    "img_input": tf.uint8,
    "det_input": tf.int8
}
# 按张量名称指定输出类型，detect签名有两个输出，分别叫output_0、output_1
converter.inference_output_type_overrides = {
    "classify": tf.float32,
    "output_0": tf.uint8,
    "output_1": tf.float32
}

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

配置时只需要对应你定义签名时给张量设置的name属性即可，不需要区分所属签名，全局唯一的张量名就能匹配到对应位置。

方案2：低版本TF手动包装签名函数实现自定义量化 #

如果你的TF版本低于2.10，没办法用原生的overrides参数，可以通过给每个签名的输入输出手动加量化/反量化算子的方式实现：

1. 先将原始模型转换为仅权重量化、输入输出保持float32的TFLite基础模型
2. 重新给每个签名的输入前插quantize算子、输出后插dequantize算子，再重新导出带自定义量化规则的多签名模型

# 第一步：导出基础量化模型（输入输出保持float32）
converter = tf.lite.TFLiteConverter.from_saved_model("./saved_model", signature_keys=["classify", "detect"])
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# 这里不指定全局输入输出类型，默认保持float32
base_tflite = converter.convert()

# 第二步：加载基础模型，手动包装每个签名添加量化反量化逻辑
interpreter = tf.lite.Interpreter(model_content=base_tflite)
signatures = interpreter.get_signature_list()

wrapped_signatures = {}
# 包装classify签名：输入uint8，输出float32
@tf.function(input_signature=[tf.TensorSpec(shape=[None,224,224,3], dtype=tf.uint8, name="img_input")])
def wrapped_classify(img_input):
    # 先反量化到float32喂给原始模型
    float_input = tf.cast(img_input, tf.float32) / 255.0
    raw_output = interpreter.get_signature_runner("classify")(img_input=float_input)["output_0"]
    return raw_output
wrapped_signatures["classify"] = wrapped_classify

# 包装detect签名：输入int8，第一个输出uint8，第二个输出float32
@tf.function(input_signature=[tf.TensorSpec(shape=[None,224,224,3], dtype=tf.int8, name="det_input")])
def wrapped_detect(det_input):
    float_input = tf.cast(det_input, tf.float32) / 127.0
    raw_outputs = interpreter.get_signature_runner("detect")(det_input=float_input)
    output_0 = tf.cast(raw_outputs["output_0"] * 255.0, tf.uint8)
    output_1 = raw_outputs["output_1"]
    return output_0, output_1
wrapped_signatures["detect"] = wrapped_detect

# 重新导出带自定义量化规则的多签名TFLite模型
converter = tf.lite.TFLiteConverter.from_concrete_functions(
    [f.get_concrete_function() for f in wrapped_signatures.values()],
)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
final_tflite = converter.convert()

验证配置是否生效 #

转换完成后可以用以下代码检查每个签名的输入输出类型是否符合预期：

interpreter = tf.lite.Interpreter(model_path="multi_sig_quant.tflite")
for sig_name in interpreter.get_signature_list():
    sig_runner = interpreter.get_signature_runner(sig_name)
    print(f"签名{sig_name}输入类型：", {k: v.dtype for k, v in sig_runner.get_input_details().items()})
    print(f"签名{sig_name}输出类型：", {k: v.dtype for k, v in sig_runner.get_output_details().items()})

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