TensorFlow 2.x加载未知架构TF Hub预训练模型求助

1. 先确认模型的可用签名标签 #

如果serving_default标签报错，先加载模型并查看所有支持的签名：

import tensorflow as tf

model_path = '/path/to/your/model/directory'
# 不带标签直接加载模型
model = tf.saved_model.load(model_path)
# 打印所有可用的签名名称
print("模型支持的签名:", list(model.signatures.keys()))

2. 获取模型的推理函数 #

根据打印出的签名名称，获取对应的推理函数，比如签名是serving_default：

infer_fn = model.signatures['serving_default']

如果有其他签名，替换成对应的key即可。

3. 查看输入输出的格式要求 #

不清楚模型架构时，先确认输入输出的张量形状和类型，避免格式不匹配报错：

# 查看输入参数信息
print("输入信息:")
for input_name, input_spec in infer_fn.structured_input_signature[1].items():
    print(f"名称: {input_name}, 形状: {input_spec.shape}, 类型: {input_spec.dtype}")

# 查看输出参数信息
print("\n输出信息:")
for output_name, output_spec in infer_fn.structured_outputs.items():
    print(f"名称: {output_name}, 形状: {output_spec.shape}, 类型: {output_spec.dtype}")

4. 准备输入并执行推理 #

根据上面得到的输入要求，准备符合格式的数据后调用推理函数。举个图像模型的示例：

# 读取并预处理测试图片（根据模型输入要求调整）
image_path = '/path/to/test/image.jpg'
image = tf.io.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize(image, (224, 224))  # 匹配模型输入尺寸
image = tf.expand_dims(image, 0)  # 增加batch维度
image = tf.cast(image, tf.float32) / 255.0  # 归一化（按模型要求调整）

# 执行推理
outputs = infer_fn(image)
# 获取输出结果（取第一个输出为例）
result = outputs[list(outputs.keys())[0]]
print("推理结果:", result)

5. TF Hub专用加载方式（可选） #

如果是标准TF Hub模块，用tfhub.load()加载可能更适配：

import tensorflow_hub as hub

model = hub.load(model_path)
# 同样查看签名和输入输出信息
print("模型支持的签名:", list(model.signatures.keys()))

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