tensorflow 加载预训练模型

最新推荐文章于 2024-09-26 12:17:07 发布

原创最新推荐文章于 2024-09-26 12:17:07 发布 · 1.8k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow

tensorflow学习专栏收录该内容

1 篇文章

订阅专栏

本文介绍了使用TensorFlow框架加载预训练模型的过程，并给出了具体的Python代码实现。通过两个实例展示了如何进行模型训练并保存，以及如何加载已保存的模型进行预测。

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

加载预训练模型需要具备两个条件：1.框架结构（知道每一层的名字），2. 预训练好的模型文件.ckpt

加载预训练模型代码如下：

import tensorflow
as tf
import numpy 
as np

weights_1 = tf.Variable(tf.zeros([3,4]))
# weights_2 = tf.Variable(tf.zeros([4,3]))

sess = tf.InteractiveSession()
saver = tf.train.Saver()

saver.restore(sess, '/tmp/checkpoint/model.ckpt')
o_test = np.array([[4.0,
3.0, 
2.0]], dtype=
'float32')
label = tf.matmul(o_test, weights_1)
# label = tf.matmul(label, weights_2)
print sess.run(label)

以上为加载模型代码，可以写全变量名，也可只写一部分。可根据输出来定。

其中，model.ckpt训练模型代码如下：

import tensorflow
as tf
import numpy 
as np

i_data = np.array([[5.0,
3.0, 
2.0]], dtype
= 'float32')
i_label= np.array([[15.0,
10.0, 
22.0]], dtype
= 'float32')

weights_1 = tf.Variable(tf.zeros([3,4]))
out_1 = tf.matmul(i_data, weights_1)

weights_2 = tf.Variable(tf.zeros([4,3]))
out = tf.matmul(out_1, weights_2)

init_op = tf.global_variables_initializer()
saver = tf.train.Saver()

loss = tf.reduce_mean(tf.square(out
- i_label))
training = tf.train.GradientDescentOptimizer(0.01).minimize(loss)

sess = tf.Session()
sess.run(init_op)
for i 
in range(20000):
sess.run(training)
save_path = saver.save(sess,
'/tmp/checkpoint/model.ckpt')