python - Tensorflow 2.0:如何在使用 tf.saved_model 时更改输出签名
问题描述
我想更改保存的模型的输入和输出签名,我使用 tf.Module 对象来构建主模型的操作。
class Generator(tf.Module):
def __init__(....):
super(Generator, self).__init__(name=name)
...
with self.name_scope:
...
@tf.Module.with_name_scope
def __call__(self, input):
...
@tf.function
def serve_function(self, input):
out = self.__call__(input)
return out
call = model.Generator.serve_function.get_concrete_function(tf.TensorSpec([None, 256, 256, 3], tf.float32))
tf.saved_model.save(model.Generator, os.path.join(train_log_dir, 'frozen'))
然后我正在加载模型,但我有签名“default_serving”和“output_0”,我该如何更改?
解决方案
我想出了一种在不使用 tf.Module 的情况下定义输出签名的方法,方法是定义 atf.function
返回输出字典,其中字典中使用的键将是输出名称。
# Create the model
model = ...
# Train the model
model.fit(...)
# Define where to save the model
export_path = "..."
@tf.function()
def my_predict(my_prediction_inputs):
inputs = {
'my_serving_input': my_prediction_inputs,
}
prediction = model(inputs)
return {"my_prediction_outputs": prediction}
my_signatures = my_predict.get_concrete_function(
my_prediction_inputs=tf.TensorSpec([None,None], dtype=tf.dtypes.float32, name="my_prediction_inputs")
)
# Save the model.
tf.saved_model.save(
model,
export_dir=export_path,
signatures=my_signatures
)
这会产生以下签名:
signature_def['serving_default']:
The given SavedModel SignatureDef contains the following input(s):
inputs['my_prediction_inputs'] tensor_info:
dtype: DT_FLOAT
shape: (-1, -1)
name: serving_default_my_prediction_inputs:0
The given SavedModel SignatureDef contains the following output(s):
outputs['my_prediction_outputs'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict
推荐阅读
- javascript - 如何加载图像并绘制图像而不会出错
- bash - 为什么在 .bat 文件 (WSL) 中运行“bash”后终端会停止以及如何正确运行它?
- excel - 如何从即时窗口调用 Excel 自定义功能区按钮?
- aws-lambda - 当我在本地部署 aws lambda 时,console.log 输出在哪里?
- python - 如何用烧瓶测试和模拟 mongodb?
- javascript - nuxt js样式作用域我不知道如何更改另一个组件的样式
- intellij-idea - 修复 JetBrains IDE 图形错误,其中菜单窗口显示在多个显示器上
- postgresql - 无法连接到 postgresql 服务器 django lightail
- java - 为什么我可以使用 mvn spring-boot:run 启动我的应用程序,而 IntelliJ run 不起作用?
- javascript - 如何设置默认标头以解决 422 状态代码错误