TensorFlow之saved_model使用笔记

最新推荐文章于 2024-03-15 23:48:28 发布

phynikesi

最新推荐文章于 2024-03-15 23:48:28 发布

阅读量799

点赞数 1

分类专栏：笔记 NLP

本文链接：https://blog.youkuaiyun.com/phynikesi/article/details/113933102

版权

笔记同时被 2 个专栏收录

9 篇文章

订阅专栏

NLP

7 篇文章

订阅专栏

这篇博客详细介绍了如何在TensorFlow中设置和使用签名定义，包括输入和输出张量的配置，以及模型的保存和加载过程。通过`tf.saved_model.builder.SavedModelBuilder`和`tf.saved_model.loader.load`进行模型的构建和加载，并展示了如何在会话中运行模型进行预测。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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)