tensorflow中模型持久化/重新加载

最新推荐文章于 2023-01-07 14:27:58 发布

原创最新推荐文章于 2023-01-07 14:27:58 发布 · 603 阅读

1 ·

CC 4.0 BY-SA版权

机器学习专栏收录该内容

26 篇文章

订阅专栏

本文介绍了如何在TensorFlow中保存和恢复模型，详细阐述了模型持久化的重要性，以及使用tf.saved_model API进行模型保存和重新加载的步骤。

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

import tensorflow as tf
# 模型持久化==================
v1 = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32),name='v1')
v2 = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32), name='v2')
result = v1+v2

saver = tf.train.Saver()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    #模型运行
    print(sess.run(result))
    for i in range(5):
        #保存
        saver.save(sess,'./models/test3/model.ckpt',global_step=i)



#模拟模型加载的代码==============

# 第一种方法---必须定义一个和模型存储的时候一样的一个计算图
v1 = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32),name='v1')
v2 = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32), name='v2')
result = v1+v2

saver = tf.train.Saver()

with tf.Session() as sess:
    saver.restore(sess,'./models/test3/model.ckpt-4')
    print(sess.run([result,v1,v2]))


#第二种方法---不定义图 直接加载模型(根据meta data的原数据信息)

saver = tf.train.import_meta_graph('./models/test3/model.ckpt-4.meta')

with tf.Session() as sess:
    saver.restore(sess,'./models/test3/model.ckpt-4')
    print(sess.run(tf.get_default_graph().get_tensor_by_name('add:0')))



# 第三种---更改映射关系
d = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32),name='a')
f = tf.Variable(initial_value=tf.random_uniform([1],0,10,dtype=tf.float32), name='b')
result = d+f

saver = tf.train.Saver(var_list={"v1":d,"v2":f})

with tf.Session() as sess:
    saver.restore(sess,'./models/test3/model.ckpt-4')
    print(sess.run([result,d,f]))