Tensorflow 更改模型保存的ckpt文件中的变量名

简介

因为需要用到多个pre-train的CNN模型，想到一个更改ckpt文件变量名的方法，这样只需在代码中更改scope名即可。

核心代码

以inception v3举例，inception_v3.ckpt下载

import tensorflow as tf

def rename_var(ckpt_path, new_ckpt_path):
    with tf.Session() as sess:
        for var_name, _ in tf.contrib.framework.list_variables(ckpt_path):
            print(var_name)
            var = tf.contrib.framework.load_variable(ckpt_path, var_name)
            new_var_name = var_name.replace('InceptionV3', 'IV')
            var = tf.Variable(var, name=new_var_name)

        saver = tf.train.Saver()
        sess.run(tf.global_variables_initializer())
        saver.save(sess, new_ckpt_path)

ckpt_path = './inception_v3.ckpt'
new_ckpt_path = './iv.ckpt'
rename_var(ckpt_path, new_ckpt_path)

这样就实现了将inception_v3.ckpt里面的变量名以InceptionV3开头的换成了IV开头的。
注：经过验证，两个模型参数一致。详情代码见tensorflow-function-examples
注：最好让新名称和旧名称差距大点，别有相同前缀，slim.assgin_from_checkpoint_fn()有bug，相同前缀的话，会出现有些变量在ckpt中找不到。

加载预训练模型

以下举个简单的例子，至于如何使用slim的预训练模型，可以自行网上找博客。这里不做赘述。

import tensorflow as tf
from tensorflow.contrib import slim

...
...
init_cnn1_fn = slim.assgin_from_checkpoint_fn(ckpt_path, slim.get_model_variables('InceptionV3'))
init_cnn2_fn = slim.assgin_from_checkpoint_fn(new_ckpt_path, slim.get_model_variables('IV'))
...
...
with tf.Session() as sess:
	sess.run(tf.global_variables_initializer())
	init_cnn1_fn(sess)
	init_cnn2_fn(sess)
	...

参考

https://stackoverflow.com/questions/37086268/rename-variable-scope-of-saved-model-in-tensorflow