SageMaker训练SSD模型转TFLite后在Swift中报错：输出数量不符

问题原因 #

你使用的tensorflow-od1-ssd-resnet50-v1-fpn-640x640-coco17-tpu-8是针对TPU优化的预训练模型，训练后导出的SavedModel包含了多尺度特征层的中间输出（共8个），但Swift的TFLite目标检测API仅兼容标准Mobile SSD的4个核心输出（检测框、类别、分数、检测数量），导致输出数量不匹配报错。

解决方案 #

通过过滤SavedModel的输出，只保留必要的4个核心输出后再转换为TFLite模型，具体步骤如下：

步骤1：查看SavedModel的签名与输出 #

先确认模型的推理签名和输出节点名称：

import tensorflow as tf

saved_model_dir = "你的SavedModel存储路径"
model = tf.saved_model.load(saved_model_dir)
# 查看可用的签名（通常为`serving_default`）
print("可用签名：", list(model.signatures.keys()))
# 查看签名对应的输出结构
infer_signature = model.signatures['serving_default']
print("输出结构：", infer_signature.structured_outputs)

输出结果中你会看到8个输出项，我们只需要保留detection_boxes、detection_classes、detection_scores、num_detections这四个。

步骤2：重新导出过滤后的SavedModel #

定义一个只返回4个核心输出的推理函数，重新导出模型：

@tf.function(input_signature=infer_signature.input_signature)
def filtered_inference(input_tensor):
    outputs = infer_signature(input_tensor)
    return {
        'detection_boxes': outputs['detection_boxes'],
        'detection_classes': outputs['detection_classes'],
        'detection_scores': outputs['detection_scores'],
        'num_detections': outputs['num_detections']
    }

# 保存过滤后的模型
new_saved_model_dir = "过滤后的SavedModel路径"
tf.saved_model.save(model, new_saved_model_dir, signatures={'serving_default': filtered_inference})

步骤3：转换为TFLite模型 #

使用过滤后的SavedModel重新转换：

converter = tf.lite.TFLiteConverter.from_saved_model(new_saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
tflite_model = converter.convert()

# 保存最终的TFLite模型
with open("目标路径/filtered_model.tflite", 'wb') as f:
    f.write(tflite_model)

简化方案（直接指定输出节点） #

如果不想重新导出模型，也可以在转换时直接指定要保留的输出数组：

converter = tf.lite.TFLiteConverter.from_saved_model(
    saved_model_dir,
    signature_keys=['serving_default'],
    output_arrays=['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections']
)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,
    tf.lite.OpsSet.SELECT_TF_OPS
]
tflite_model = converter.convert()

with open("目标路径/filtered_model.tflite", 'wb') as f:
    f.write(tflite_model)

用上述方法得到的TFLite模型仅包含4个必要输出，即可在Swift应用中正常进行目标检测。

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