tensorflow 加载预训练模型

加载预训练模型需要具备两个条件：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')