signature设置

x1 = tf.placeholder(tf.int32, shape=[None, None], name='x1')
x2 = tf.placeholder(tf.int32, shape=[None, None], name='x2')
……
y = output_tensor
loss = loss_tensor

inputs = {
            'x1': tf.saved_model.utils.build_tensor_info(x1),
            'x2': tf.saved_model.utils.build_tensor_info(x2)
}
outputs = {
            'y': tf.saved_model.utils.build_tensor_info(y),
            'loss': tf.saved_model.utils.build_tensor_info(loss)
}

signature = tf.saved_model.signature_def_utils.build_signature_def(
inputs=inputs, outputs=outputs, method_name='sig_cpu')  # 名字自定义

模型保存

# 1、建立builder
# 2、保存会话（主要是计算图和参数），有signature也一并保存
# 序列化写入磁盘

with tf.Session() as sess:
    model_path = path + '/model_save'
    bulder = tf.saved_model.builder.SavedModelBuilder(model_path)
    sig_config = {'save_sig': signature}
    builder.add_meta_graph_and_variables(sess, ['cpu_server'], sig_config)
    builder.save()

模型加载与使用

with tf.Session() as sess:
    # 不用执行初始化
    meta_graph_def = tf.saved_model.loader.load(sess, ['cpu_server'],  
                       model_path+'/model_save')
    signature = meta_graph_def.signature_def
    signature_key = 'save_sig'
    input_x1 = 'x1'
    input_x2 = 'x2'
    output_y = 'y'
    out_loss = 'loss'
    
    x1 = signature[signature_key].inputs[input_x1].name
    x2 = signature[signature_key].inputs[input_x2].name
    y = signature[signature_key].inputs[output_y].name
    loss = signature[signature_key].outputs[output_loss].name

    feed = {x1: , x2: , y: , loss: }
    res, _ = sess.run([y, loss], feed_dict=feed)