使用自定义签名defs保存TF2 keras模型

时间:2019-06-19 04:13:19

标签: python tensorflow keras tensorflow2.0

我有一个Keras(顺序)模型,可以使用Tensorflow 1.13中的自定义签名defs进行保存,如下所示:

from tensorflow.saved_model.utils import build_tensor_info
from tensorflow.saved_model.signature_def_utils import predict_signature_def, build_signature_def

model = Sequential() // with some layers

builder = tf.saved_model.builder.SavedModelBuilder(export_path)

score_signature = predict_signature_def(
    inputs={'waveform': model.input},
    outputs={'scores': model.output})

metadata = build_signature_def(
    outputs={'other_variable': build_tensor_info(tf.constant(1234, dtype=tf.int64))})

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  builder.add_meta_graph_and_variables(
      sess=sess,
      tags=[tf.saved_model.tag_constants.SERVING],
      signature_def_map={'score': score_signature, 'metadata': metadata})
  builder.save()

将模型迁移到TF2 keras很酷:),但我不知道如何保存具有上述相同签名的模型。我应该使用新的tf.saved_model.save()还是tf.keras.experimental.export_saved_model()?上面的代码应如何用TF2编写?

关键要求:

  • 该模型具有得分签名和元数据签名
  • 元数据签名包含1个或多个常量

1 个答案:

答案 0 :(得分:0)

解决方案是使用每个签名定义的函数创建一个tf.Module

class MyModule(tf.Module):
  def __init__(self, model, other_variable):
    self.model = model
    self._other_variable = other_variable

  @tf.function(input_signature=[tf.TensorSpec(shape=(None, None, 1), dtype=tf.float32)])
  def score(self, waveform):
    result = self.model(waveform)
    return { "scores": results }

  @tf.function(input_signature=[])
  def metadata(self):
    return { "other_variable": self._other_variable }

然后保存模块(不是模型):

module = MyModule(model, 1234)
tf.saved_model.save(module, export_path, signatures={ "score": module.score, "metadata": module.metadata })

在TF2上使用Keras模型进行了测试。