2016-05-07 78 views
11

是否可以重新命名tensorflow中给定模型的变量范围?在TensorFlow中重命名已保存模型的变量范围

例如,我创建了一个逻辑回归模型MNIST数字,根据教程:

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]) 

... # some training happens 

saver.save(sess, 'my-model') 

现在我要重新加载保存的模型在'my-first-scope'变量范围,然后再保存一切新文件并在'my-second-scope'的新变量范围内。

回答

7

您可以使用tf.contrib.framework.list_variablestf.contrib.framework.load_variable如下,以实现自己的目标:

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') 
+0

这要求您已经使用先前的作用域名称构造了图形和所有内容,因为要恢复您需要定义该图形的检查点。 如果您只有检查点文件,您可以替换其中的范围名称吗? – npit

18

基于keveman的回答,我创建了一个python脚本,就可以执行重命名任何TensorFlow检查点的变量:

https://gist.github.com/batzner/7c24802dd9c5e15870b4b56e22135c96

您可以替换变量名称中的子字符串并为所有名称添加前缀。呼叫与

python tensorflow_rename_variables.py --checkpoint_dir=path/to/dir 

脚本使用可选参数

--replace_from=substr --replace_to=substr --add_prefix=abc --dry_run 

这里是脚本的核心功能:

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) 

例子:

python tensorflow_rename_variables.py --checkpoint_dir=path/to/dir --replace_from=scope1 --replace_to=scope1/model --add_prefix=abc/ 

将重命名变量scope1/Variable1abc/scope1/model/Variable1

+0

我得到这个脚本的错误:ValueError:找不到'checkpoint'文件或检查点在给定的目录./fi – ryuzakinho

+1

@ryuzakinho,你需要指定一个包含'checkpoint'文件的目录。有关更多信息,请参阅https://www.tensorflow.org/programmers_guide/variables#checkpoint_files。 –

+1

实际上,出于某种原因,我有write_state = False。因此,它不会创建检查点文件。 – ryuzakinho