Rename variable scope of saved model in TensorFlow
You can use tf.contrib.framework.list_variables
and tf.contrib.framework.load_variable
as follows to achieve your goal :
with tf.Graph().as_default(), tf.Session().as_default() as sess:
with tf.variable_scope('my-first-scope'):
NUM_IMAGE_PIXELS = 784
NUM_CLASS_BINS = 10
x = tf.placeholder(tf.float32, shape=[None, NUM_IMAGE_PIXELS])
y_ = tf.placeholder(tf.float32, shape=[None, NUM_CLASS_BINS])
W = tf.Variable(tf.zeros([NUM_IMAGE_PIXELS,NUM_CLASS_BINS]))
b = tf.Variable(tf.zeros([NUM_CLASS_BINS]))
y = tf.nn.softmax(tf.matmul(x,W) + b)
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
saver = tf.train.Saver([W, b])
sess.run(tf.global_variables_initializer())
saver.save(sess, 'my-model')
vars = tf.contrib.framework.list_variables('.')
with tf.Graph().as_default(), tf.Session().as_default() as sess:
new_vars = []
for name, shape in vars:
v = tf.contrib.framework.load_variable('.', name)
new_vars.append(tf.Variable(v, name=name.replace('my-first-scope', 'my-second-scope')))
saver = tf.train.Saver(new_vars)
sess.run(tf.global_variables_initializer())
saver.save(sess, 'my-new-model')
Based on keveman's answer, I created a python script, which you can execute to rename the variables of any TensorFlow checkpoint:
https://gist.github.com/batzner/7c24802dd9c5e15870b4b56e22135c96
You can replace substrings in the variable names and add a prefix to all names. Call the script with
python tensorflow_rename_variables.py --checkpoint_dir=path/to/dir
with the optional arguments
--replace_from=substr --replace_to=substr --add_prefix=abc --dry_run
Here is the script's core function:
def rename(checkpoint_dir, replace_from, replace_to, add_prefix, dry_run=False):
checkpoint = tf.train.get_checkpoint_state(checkpoint_dir)
with tf.Session() as sess:
for var_name, _ in tf.contrib.framework.list_variables(checkpoint_dir):
# Load the variable
var = tf.contrib.framework.load_variable(checkpoint_dir, var_name)
# Set the new name
new_name = var_name
if None not in [replace_from, replace_to]:
new_name = new_name.replace(replace_from, replace_to)
if add_prefix:
new_name = add_prefix + new_name
if dry_run:
print('%s would be renamed to %s.' % (var_name, new_name))
else:
print('Renaming %s to %s.' % (var_name, new_name))
# Rename the variable
var = tf.Variable(var, name=new_name)
if not dry_run:
# Save the variables
saver = tf.train.Saver()
sess.run(tf.global_variables_initializer())
saver.save(sess, checkpoint.model_checkpoint_path)
Example:
python tensorflow_rename_variables.py --checkpoint_dir=path/to/dir --replace_from=scope1 --replace_to=scope1/model --add_prefix=abc/
will rename the variable scope1/Variable1
to abc/scope1/model/Variable1
.