V2EX = way to explore
V2EX 是一个关于分享和探索的地方
现在注册
已注册用户请  登录
fendouai_com
V2EX  ›  TensorFlow

TensorFlow 训练好模型参数的保存和恢复代码

  •  
  •   fendouai_com · 2017-07-16 19:48:49 +08:00 · 2997 次点击
    这是一个创建于 2697 天前的主题,其中的信息可能已经有所发展或是发生改变。

    TensorFlow 训练好模型参数的保存和恢复代码,之前就在想模型不应该每次要个结果都要重新训练一遍吧,应该训练一次就可以一直使用吧。

    TensorFlow 提供了 Saver 类,可以进行保存和恢复。下面是 TensorFlow-Examples 项目中提供的保存和恢复代码。

    Save and Restore a model using TensorFlow.
    This example is using the MNIST database of handwritten digits
    ( http://yann.lecun.com/exdb/mnist/)
    
    Author: Aymeric Damien
    Project: https://github.com/aymericdamien/TensorFlow-Examples/
    '''
    
    from __future__ import print_function
    
    # Import MNIST data
    from tensorflow.examples.tutorials.mnist import input_data
    mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
    
    import tensorflow as tf
    
    # Parameters
    learning_rate = 0.001
    batch_size = 100
    display_step = 1
    model_path = "/tmp/model.ckpt"
    
    # Network Parameters
    n_hidden_1 = 256 # 1st layer number of features
    n_hidden_2 = 256 # 2nd layer number of features
    n_input = 784 # MNIST data input (img shape: 28*28)
    n_classes = 10 # MNIST total classes (0-9 digits)
    
    # tf Graph input
    x = tf.placeholder("float", [None, n_input])
    y = tf.placeholder("float", [None, n_classes])
    
    
    # Create model
    def multilayer_perceptron(x, weights, biases):
        # Hidden layer with RELU activation
        layer_1 = tf.add(tf.matmul(x, weights['h1']), biases['b1'])
        layer_1 = tf.nn.relu(layer_1)
        # Hidden layer with RELU activation
        layer_2 = tf.add(tf.matmul(layer_1, weights['h2']), biases['b2'])
        layer_2 = tf.nn.relu(layer_2)
        # Output layer with linear activation
        out_layer = tf.matmul(layer_2, weights['out']) + biases['out']
        return out_layer
    
    # Store layers weight & bias
    weights = {
        'h1': tf.Variable(tf.random_normal([n_input, n_hidden_1])),
        'h2': tf.Variable(tf.random_normal([n_hidden_1, n_hidden_2])),
        'out': tf.Variable(tf.random_normal([n_hidden_2, n_classes]))
    }
    biases = {
        'b1': tf.Variable(tf.random_normal([n_hidden_1])),
        'b2': tf.Variable(tf.random_normal([n_hidden_2])),
        'out': tf.Variable(tf.random_normal([n_classes]))
    }
    
    # Construct model
    pred = multilayer_perceptron(x, weights, biases)
    
    # Define loss and optimizer
    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=y))
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost)
    
    # Initializing the variables
    init = tf.global_variables_initializer()
    
    # 'Saver' op to save and restore all the variables
    saver = tf.train.Saver()
    
    # Running first session
    print("Starting 1st session...")
    with tf.Session() as sess:
        # Initialize variables
        sess.run(init)
    
        # Training cycle
        for epoch in range(3):
            avg_cost = 0.
            total_batch = int(mnist.train.num_examples/batch_size)
            # Loop over all batches
            for i in range(total_batch):
                batch_x, batch_y = mnist.train.next_batch(batch_size)
                # Run optimization op (backprop) and cost op (to get loss value)
                _, c = sess.run([optimizer, cost], feed_dict={x: batch_x,
                                                              y: batch_y})
                # Compute average loss
                avg_cost += c / total_batch
            # Display logs per epoch step
            if epoch % display_step == 0:
                print("Epoch:", '%04d' % (epoch+1), "cost=", \
                    "{:.9f}".format(avg_cost))
        print("First Optimization Finished!")
    
        # Test model
        correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
        # Calculate accuracy
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
        print("Accuracy:", accuracy.eval({x: mnist.test.images, y: mnist.test.labels}))
    
        # Save model weights to disk
        save_path = saver.save(sess, model_path)
        print("Model saved in file: %s" % save_path)
    
    # Running a new session
    print("Starting 2nd session...")
    with tf.Session() as sess:
        # Initialize variables
        sess.run(init)
    
        # Restore model weights from previously saved model
        saver.restore(sess, model_path)
        print("Model restored from file: %s" % save_path)
    
        # Resume training
        for epoch in range(7):
            avg_cost = 0.
            total_batch = int(mnist.train.num_examples / batch_size)
            # Loop over all batches
            for i in range(total_batch):
                batch_x, batch_y = mnist.train.next_batch(batch_size)
                # Run optimization op (backprop) and cost op (to get loss value)
                _, c = sess.run([optimizer, cost], feed_dict={x: batch_x,
                                                              y: batch_y})
                # Compute average loss
                avg_cost += c / total_batch
            # Display logs per epoch step
            if epoch % display_step == 0:
                print("Epoch:", '%04d' % (epoch + 1), "cost=", \
                    "{:.9f}".format(avg_cost))
        print("Second Optimization Finished!")
    
        # Test model
        correct_prediction = tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1))
        # Calculate accuracy
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
        print("Accuracy:", accuracy.eval(
            {x: mnist.test.images, y: mnist.test.labels}))
    

    博客: http://www.tensorflownews.com/2017/07/16/save-and-restore-a-model-using-tensorflow/

    5 条回复    2017-07-16 21:01:18 +08:00
    snnn
        1
    snnn  
       2017-07-16 20:22:42 +08:00   ❤️ 1
    啥乱七八糟的啊……
    init 和 restore 二选一就够了
    fendouai_com
        2
    fendouai_com  
    OP
       2017-07-16 20:38:10 +08:00
    @snnn init 是初始化,restore 是恢复保存的数据,这两个功能不一样吧?我也刚学,多谢大佬指教。
    snnn
        3
    snnn  
       2017-07-16 20:43:17 +08:00   ❤️ 1
    @fendouai_com 一般来说,执行了 init 就没必要执行 restore,执行了 restore 就没必要执行 init。它们都是为了初始化变量用的。
    fendouai_com
        4
    fendouai_com  
    OP
       2017-07-16 21:00:39 +08:00
    @snnn 多谢大佬指教,刚才修改了一下,确实是重复了。
    fendouai_com
        5
    fendouai_com  
    OP
       2017-07-16 21:01:18 +08:00
    # Running a new session
    print("Starting 2nd session...")
    with tf.Session() as sess:
    # Initialize variables
    # 下面这个 init 可以注释掉
    sess.run(init)
    关于   ·   帮助文档   ·   博客   ·   API   ·   FAQ   ·   实用小工具   ·   1537 人在线   最高记录 6679   ·     Select Language
    创意工作者们的社区
    World is powered by solitude
    VERSION: 3.9.8.5 · 18ms · UTC 17:09 · PVG 01:09 · LAX 09:09 · JFK 12:09
    Developed with CodeLauncher
    ♥ Do have faith in what you're doing.