tensorflow模型保存、加载之变量重命名

左手的ㄟ右手 2022-02-25 01:47 349阅读 0赞

话不多说，干就完了。

--------------------

**变量重命名的用处？**

**简单定义：**简单来说就是将模型A中的参数parameter\_A赋给模型B中的parameter\_B

**使用场景：**当需要使用已经训练好的模型参数，尤其是使用别人训练好的模型参数时，往往别人模型中的参数命名方式与自己当前的命名方式不同，所以在加载模型参数时需要对参数进行重命名，使得代码更简洁易懂。

**实现方法：**

1）、模型保存

import os
    import tensorflow as tf
    
    weights = tf.Variable(initial_value=tf.truncated_normal(shape=[1024, 2],
                                                            mean=0.0,
                                                            stddev=0.1),
                          dtype=tf.float32,
                          name="weights")
    biases = tf.Variable(initial_value=tf.zeros(shape=[2]),
                         dtype=tf.float32,
                         name="biases")
    
    weights_2 = tf.Variable(initial_value=weights.initialized_value(),
                            dtype=tf.float32,
                            name="weights_2")
    
    # saver checkpoint
    if os.path.exists("checkpoints") is False:
        os.makedirs("checkpoints")
    
    saver = tf.train.Saver()
    with tf.Session() as sess:
        init_op = [tf.global_variables_initializer()]
        sess.run(init_op)
        saver.save(sess=sess, save_path="checkpoints/variable.ckpt")

2）、模型加载（变量名称保持不变）

import tensorflow as tf
    from matplotlib import pyplot as plt
    import os
    
    current_path = os.path.dirname(os.path.abspath(__file__))
    
    def restore_variable(sess):
        # need not initilize variable, but need to define the same variable like checkpoint
        weights = tf.Variable(initial_value=tf.truncated_normal(shape=[1024, 2],
                                                                mean=0.0,
                                                                stddev=0.1),
                              dtype=tf.float32,
                              name="weights")
        biases = tf.Variable(initial_value=tf.zeros(shape=[2]),
                             dtype=tf.float32,
                             name="biases")
    
        weights_2 = tf.Variable(initial_value=weights.initialized_value(),
                                dtype=tf.float32,
                                name="weights_2")
    
        saver = tf.train.Saver()
    
        ckpt_path = os.path.join(current_path, "checkpoints", "variable.ckpt")
        saver.restore(sess=sess, save_path=ckpt_path)
    
        weights_val, weights_2_val = sess.run(
            [
                tf.reshape(weights, shape=[2048]),
                tf.reshape(weights_2, shape=[2048])
            ]
        )
    
        plt.subplot(1, 2, 1)
        plt.scatter([i for i in range(len(weights_val))], weights_val)
        plt.subplot(1, 2, 2)
        plt.scatter([i for i in range(len(weights_2_val))], weights_2_val)
        plt.show()
    
    
    if __name__ == '__main__':
        with tf.Session() as sess:
            restore_variable(sess)

3）、模型加载（变量重命名）

import tensorflow as tf
    from matplotlib import pyplot as plt
    import os
    
    current_path = os.path.dirname(os.path.abspath(__file__))
    
    
    def restore_variable_renamed(sess):
        conv1_w = tf.Variable(initial_value=tf.truncated_normal(shape=[1024, 2],
                                                                mean=0.0,
                                                                stddev=0.1),
                              dtype=tf.float32,
                              name="conv1_w")
        conv1_b = tf.Variable(initial_value=tf.zeros(shape=[2]),
                              dtype=tf.float32,
                              name="conv1_b")
    
        conv2_w = tf.Variable(initial_value=conv1_w.initialized_value(),
                              dtype=tf.float32,
                              name="conv2_w")
    
        # variable named 'weights' in ckpt assigned to current variable conv1_w
        # variable named 'biases' in ckpt assigned to current variable conv1_b
        # variable named 'weights_2' in ckpt assigned to current variable conv2_w
        saver = tf.train.Saver({
            "weights": conv1_w,
            "biases": conv1_b,
            "weights_2": conv2_w
        })
    
        ckpt_path = os.path.join(current_path, "checkpoints", "variable.ckpt")
        saver.restore(sess=sess, save_path=ckpt_path)
    
        conv1_w__val, conv2_w__val = sess.run(
            [
                tf.reshape(conv1_w, shape=[2048]),
                tf.reshape(conv2_w, shape=[2048])
            ]
        )
    
        plt.subplot(1, 2, 1)
        plt.scatter([i for i in range(len(conv1_w__val))], conv1_w__val)
        plt.subplot(1, 2, 2)
        plt.scatter([i for i in range(len(conv2_w__val))], conv2_w__val)
        plt.show()
    
    
    if __name__ == '__main__':
        with tf.Session() as sess:
            restore_variable_renamed(sess)

**总结：**

**    \#** **之前模型中叫 'weights'的变量赋值给当前的conv1\_w变量**

**    \#** **之前模型中叫 'biases' 的变量赋值给当前的conv1\_b变量**

**    \#** **之前模型中叫 'weights\_2'的变量赋值给当前的conv2\_w变量**

saver = tf.train.Saver(**\{**

**        "weights": conv1\_w,**

**        "biases": conv1\_b,**

**        "weights\_2": conv2\_w**

**    \}**)