tensorflow模型的保存和重载

川长思鸟来 2022-01-27 04:31 292阅读 0赞

# tensorflow提供了tf.train.Saver()函数来实现模型的保存和重载 #

### 模型的保存 ###

import tensorflow as tf
    import numpy as np
    
    # 创建两个变量
    W = tf.Variable([[1, 2, 3], [3, 2, 1]], dtype=tf.float32, name='weights')
    b = tf.Variable([1, 2, 3], dtype=tf.float32, name='weights')
    
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        saver.save(sess, './finetuning/model1.ckpt', global_step = 100)
        print('模型保存成功')

### 模型的重新加载 ###

import tensorflow as tf
    import numpy as np
    
    # 创建两个变量，相当于给上面创建的两个变量提供存放的地点， 本身没有特别的意义
    # 具体定义的值不重要，只要定义的shape和dtype一致就可以了
    W = tf.Variable(np.zeros(6).reshape((2, 3)), dtype=tf.float32, name='weights')
    b = tf.Variable([1, 1, 1], dtype=tf.float32, name='weights')
    
    # 记住这里不需要对变量进行初始化
    saver = tf.train.Saver()
    with tf.Session() as sess:
        # 这个函数会默认找到最新更新的模型保存文件
        ckpt = tf.train.get_checkpoint_state('./finetuning/')
        saver.restore(sess, ckpt.model_checkpoint_path)
        print(sess.run(W))
        print(sess.run(b))

输出结果：

INFO:tensorflow:Restoring parameters from ./finetuning/model1.ckpt-100
    [[1. 2. 3.]
     [3. 2. 1.]]
    [1. 2. 3.]

可以看到restore的时候无论定义什么都不会输出，只是提供了一个载体而已；要注意的是：如果我们定义的是一个网络结构，那么我们还需要在restore的时候重新定义网络结构

## 自己简单测试 ##

# restore
    
    import tensorflow as tf
    from tensorflow.examples.tutorials.mnist import input_data
     
     
    # MNIST数据集相关的常数
    INPUT_NODE = 784
    OUTPUT_NODE = 10
    LAYER1_NODE = 500
    BATCH_SIZE = 100
    LEARNING_RATE_BASE = 0.8
    LEARNING_RATE_DECAY = 0.99
    REGULARIZATION_RATE = 0.0001  # 描述模型复杂度的正则化在损失函数的系数
    TRAINING_STEPS = 10000
    MOVING_AVERAGE_DECAY = 0.99   # 滑动平均衰减率
    
    # 一个辅助函数，给定神经网络的输入和所有参数，计算神经网络中的前向传播结果
    def inference(input_tensor, avg_class, weights1, biases1, weights2, biases2):
        # 当没有滑动平均类，直接使用参数当前的取值
        if avg_class == None:
            layer1 = tf.nn.relu(tf.matmul(input_tensor, weights1) + biases1)
            pred = tf.matmul(layer1, weights2) + biases2
        else:
            layer1 = tf.nn.relu(tf.matmul(input_tensor, avg_class.average(weights1)) + avg_class.average(biases1))
            pred = tf.matmul(layer1, avg_class.average(weights2)) + avg_class.average(biases2)
        return pred
     
    # 训练模型的过程
    def train(mnist):
        x = tf.placeholder(tf.float32, [None, INPUT_NODE], name='x_input')
        y = tf.placeholder(tf.float32, [None, OUTPUT_NODE], name='y_output')
        weights1 = tf.Variable(tf.truncated_normal([INPUT_NODE, LAYER1_NODE], stddev=0.1))
        biases1 = tf.Variable(tf.constant(0.1, shape=[LAYER1_NODE]))
        weights2 = tf.Variable(tf.truncated_normal([LAYER1_NODE, OUTPUT_NODE], stddev=0.1))
        biases2 = tf.Variable(tf.constant(0.1, shape=[OUTPUT_NODE]))
        
        # 计算在当前参数下神经网络前向传播的结果
        pred = inference(x, None, weights1, biases1, weights2, biases2)
        # 计算使用了滑动平均之后的前向传播结果
        global_step = tf.Variable(0, trainable=False)
        variable_average = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, global_step)
        # tf.trainable_variables函数返回的就是图上的所有变量的集合
        average_op = variable_average.apply(tf.trainable_variables())
        pred1 = inference(x, variable_average, weights1, biases1, weights2, biases2)
     
        correct_prediction = tf.equal(tf.argmax(pred1, 1), tf.argmax(y, 1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
        
        saver = tf.train.Saver()
        # 初始会话并开始训练
        with tf.Session() as sess:
            # tf.global_variables_initializer().run()
            # 准备验证数据
            # validate_feed = {x:mnist.validation.images, y:mnist.validation.labels}
            # 准备测试数据
            test_feed = {x:mnist.test.images, y:mnist.test.labels}
            ckpt = tf.train.get_checkpoint_state('./')
            saver.restore(sess, ckpt.model_checkpoint_path)
            print(sess.run(weights1))
    
            # 在训练结束之后，在测试数据上检测神经网络模型的最终正确率
            test_acc = sess.run(accuracy, feed_dict=test_feed)
            print('模型加载成功')
            print('test_acc:', test_acc)
            
    # 主程序入口
    def main(argv=None):
        # 声明处理MNIST数据集的类，这个类在初始化的时候会自动下载数据
        mnist = input_data.read_data_sets('/path/to/MNIST_DATA/', one_hot=True)
        train(mnist)
        
     
    if __name__ == '__main__':
        tf.app.run()

完整的训练流程在tensorflow基础概念博客的第四篇（好像是），上面简单修改了一下

和原始训练代码的区别有以下几点：

*  整个前向传播过程是不用进行修改的
 *  保留和输出相关的tensor的定义
 *  变量样要定义，只不过不需要进行初始化