TensorFlow修改预训练网络模型参数名称

一、链接

1、tensorflow 保存训练模型ckpt 查看ckpt文件中的变量名和对应值

https://www.cnblogs.com/adong7639/p/7764769.html

2、Tensorflow修改已训练模型变量名字的方法

https://zhuanlan.zhihu.com/p/36982683

3、如何修改在TensorFlow框架下训练保存的模型参数名称

https://blog.csdn.net/jiongnima/article/details/86632517

二、原理介绍：

       在解释代码之前，先介绍一下要用到的两个重要的接口：

1. tf.contrib.framework.list_variables。将已保存参数的（名称，形状）以列表的形式返回。在更新的TensorFlow版本中，该接口已经被整合到了tf.train.list_variables里面。

2. tf.contrib.framework.load_variable。可以传入名称，返回读取的已保存参数的值。在更新的TensorFlow版本中，该接口已经被整合到了tf.train.load_variable里面。

     在修改保存的参数名称时，做法分为以下6步：

1. 使用list_variables函数逐个读出已保存的参数名称

2. 使用load_variable函数逐个读取已保存的参数值

3. 逐个修改参数名称

4. 使用已修改的参数名称，结合tf.Variable函数逐个重建参数

5. 将已重建的参数逐个加入新参数列表

6. 使用tf.train.Saver().save将新参数列表写入硬盘

       下面放出笔者的代码，在代码中，笔者给DeepLab V2预训练的模型参数全加上了前缀“deeplab_v2”。在这里笔者使用的还是许久之前的DeepLab预训练模型，参数保存还是一个ckpt文件（deeplab_resnet.ckpt）。代码如下：

import tensorflow as tf

import argparse

import os

parser = argparse.ArgumentParser(description='')

parser.add_argument("--checkpoint_path", default='../deeplab_resnet/deeplab_resnet.ckpt', help="restore ckpt") #原参数路径

parser.add_argument("--new_checkpoint_path", default='../deeplab_resnet_altered/', help="path_for_new ckpt") #新参数保存路径

parser.add_argument("--add_prefix", default='deeplab_v2/', help="prefix for addition") #新参数名称中加入的前缀名

args = parser.parse_args()

def main():

        if not os.path.exists(args.new_checkpoint_path):

                os.makedirs(args.new_checkpoint_path)

       with tf.Session() as sess:

               new_var_list=[] #新建一个空列表存储更新后的Variable变量

              for var_name, _ in tf.contrib.framework.list_variables(args.checkpoint_path): #得到checkpoint文件中所有的参数（名字，形状）元组

                       var = tf.contrib.framework.load_variable(args.checkpoint_path, var_name) #得到上述参数的值

                       new_name = var_name

                       new_name = args.add_prefix + new_name #在这里加入了名称前缀，大家可以自由地作修改

                        #除了修改参数名称，还可以修改参数值（var）

                        print('Renaming %s to %s.' % (var_name, new_name))

                        renamed_var = tf.Variable(var, name=new_name) #使用加入前缀的新名称重新构造了参数

                       new_var_list.append(renamed_var) #把赋予新名称的参数加入空列表

              print('starting to write new checkpoint !')

               saver = tf.train.Saver(var_list=new_var_list) #构造一个保存器

               sess.run(tf.global_variables_initializer()) #初始化一下参数（这一步必做）

               model_name = 'deeplab_resnet_altered' #构造一个保存的模型名称

               checkpoint_path = os.path.join(args.new_checkpoint_path, model_name) #构造一下保存路径

               saver.save(sess, checkpoint_path) #直接进行保存

               print("done !")

if __name__ == '__main__':

    main()

在终端下面运行一下代码：

可以看到参数名称都被重置了，加上了前缀“deeplab_v2”：

在代码中设定的保存文件夹下，能够查看已保存的新参数名称的模型参数：

       由于后来的TensorFlow框架在保存模型时已经放弃了保存单个ckpt文件的做法，因此都是得到4个文件，如上所示。然后我们就可以在代码中愉快地使用新参数名称的模型进行初始化啦~

loader = tf.train.Saver(var_list=restore_vars) #设置一下要初始化哪些参数

checkpoint = tf.train.latest_checkpoint(args.checkpoint_path) #保存的新参数名的模型路径

loader.restore(sess, ckpt_path) #初始化模型参数

       到这里，本篇博文就接近尾声了。本篇博文主要讲述了如何修改TensorFlow框架下训练的参数名称，核心还是找出参数名->更改参数名->重建参数->保存。