Tensorflow2.x TF2.x 模型導出 自定 輸入層名稱 輸出層名稱

 Tensorflow2.x   TF2.x 模型導出   自定 輸入層名稱  輸出層名稱

import tensorflow as tf
import os
class MyModel(tf.keras.Model):
    def __init__(self, **kwargs):
        super(MyModel, self).__init__(**kwargs)
        self.dense_1 = tf.keras.layers.Dense(10)
        self.dense_2 = tf.keras.layers.Dense(1)
        self.dropout = tf.keras.layers.Dropout(0.1)

    def call(self, inputs, training=None):
        a, b = inputs
        hidden = self.dense_1(a + b)
        hidden = self.dropout(hidden, training=training)
        logits = self.dense_2(hidden)
        return logits

    @tf.function(input_signature=[(tf.TensorSpec([None, 10], name='a', dtype=tf.float32),
                                   tf.TensorSpec([None, 10], name='b', dtype=tf.float32))])
    def sever(self, inputs):
        return {"score": self.call(inputs, training=False)}

if __name__ == '__main__':
    model = MyModel()
    data = tf.ones((2, 10))
    out = model((data, data), training=True)
    print(out)
    model.save("SavedModel", signatures={"serving_default": model.sever})

这里自定义了sever函数对call进行包装并对server指定 input_signature ，

同时在输出时通过dic 指定输出变量名 "score" 。

最后 ，在model.save导出模型时, 指定对serving_default 签名对应的 调用函数进行。

最后通过  saved_model_cli show --all --dir SavedModel      命令对保存的模型进行查看，

发现这次自定义了 输入input 而且 还自定义 输出 output。

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['a'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 10)
        name: serving_default_a:0
    inputs['b'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 10)
        name: serving_default_b:0
 

 The given SavedModel SignatureDef contains the following output(s):
    outputs['score'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict