Search code examples
tensorflowneural-networkrestoremnistpre-trained-model

How to use a pretrained model with Tensorflow?


I know that the following is an already answered question, but even though i tried and tried all the proposed solutions, none of them solved my problem. I made this net for training over MNIST dataset. At the beginning it was deeper, but in order to focus on the problem i simplified it.

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()

Thanks everybody for the answers!

Regards,

Silvio

UPDATE

  • I tried both

    saver.restore(sess, model_to_restore)
    

    and

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

    but in both cases i had this error from terminal:

    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)]]
    

Solution

  • I think your location to the model might be wrong and i would advise you to give the following workflow a try.

    Since the saved models comprise several files i usually save them to a folder after training:

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

    This will also tell you the exact location where it has been saved.

    Then i can start my predictor inside the saved location (cd into myMNIST) and restore the model by:

    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!")