如何在Tensorflow中使用预训练模型?

时间:2017-06-02 08:19:46

标签: tensorflow neural-network restore mnist pre-trained-model

我知道以下是一个已经回答的问题,但即使我尝试并尝试了所有提出的解决方案,但它们都没有解决我的问题。 我用这个网进行了对MNIST数据集的训练。一开始它更深,但为了专注于问题,我简化了它。

mnist = mnist_data.read_data_sets('MNIST_data', one_hot=True)

# train the net
def train():
    for i in range(1000):
        batch_xs, batch_ys = mnist.train.next_batch(100)
        sess.run(train_step, feed_dict={x: batch_xs, y_: batch_ys})
        print("accuracy", sess.run(accuracy, feed_dict={x: mnist.test.images, y_: mnist.test.labels}))
        if i%100==0:
            save_path = saver.save(sess, "./tmp/model.ckpt", global_step = i, write_meta_graph=True)    
            print("Model saved in file: %s" % save_path)

# evaluate the net
def test(image, label):
    true_value = tf.argmax(label, 1)
    prediction = tf.argmax(y, 1)
    print("true value:", sess.run(true_value))
    print("predictions", sess.run(prediction, feed_dict={x:image}))

sess = tf.InteractiveSession()

x = tf.placeholder("float", shape=[None, 784])
W = tf.Variable(tf.zeros([784,10]), name = "W1")
b = tf.Variable(tf.zeros([10]), name = "B1")
y = tf.nn.softmax(tf.matmul(x,W) + b, name ="Y")
y_ = tf.placeholder("float", shape=[None, 10])
cross_entropy = -tf.reduce_sum(y_*tf.log(y))
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)
correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))

saver = tf.train.Saver()
model_to_restore="./tmp/model.ckpt-100.meta"
if os.path.isfile(model_to_restore):
    #what i have to do here?????#
else:
#this part works!#
    print("Model does not exist: training")
    train()

感谢大家的答案!

问候,

西尔维奥

更新

  • 我试过了两次

    saver.restore(sess, model_to_restore)
    

    saver = tf.train.import_meta_graph(model_to_restore)
    saver.restore(sess, model_to_restore)
    

    但是在这两种情况下我从终端都有这个错误:

    DataLossError (see above for traceback): Unable to open table     file ./tmp/model.ckpt.meta: Data loss: not an sstable (bad magic number): perhaps your file is in a different file format and you need to use a different restore operator?
     [[Node: save/RestoreV2 = RestoreV2[dtypes=[DT_FLOAT], _device="/job:localhost/replica:0/task:0/cpu:0"](_recv_save/Const_0, save/RestoreV2/tensor_names, save/RestoreV2/shape_and_slices)]]
    

1 个答案:

答案 0 :(得分:0)

我认为您对模型的位置可能有误,我建议您尝试使用以下工作流程。

由于保存的模型包含多个文件,我通常会在训练后将它们保存到文件夹中:

modelPath = "myMNIST/model"
saved_path = saver.save(sess, os.path.join(modelPath, "model.ckpt"))
print("Model saved in file: ", saved_path)

这也会告诉您保存的确切位置。

然后我可以在保存的位置内启动我的预测器(cd进入myMNIST)并通过以下方式恢复模型:

ckpt = tf.train.get_checkpoint_state("./model")
if ckpt and ckpt.model_checkpoint_path:
    print("Restored Model")
    saver.restore(sess, ckpt.model_checkpoint_path)
else:
    print("Could not restore model!")