【TensorFlow】使用slim从ckpt里导出指定层的参数

清疚 2024-04-18 14:49 70阅读 0赞

利用函数 `slim.assgin_from_checkpoint`

相关文档：[https://github.com/tensorflow/tensorflow/blob/129665119ea60640f7ed921f36db9b5c23455224/tensorflow/contrib/slim/python/slim/learning.py][https_github.com_tensorflow_tensorflow_blob_129665119ea60640f7ed921f36db9b5c23455224_tensorflow_contrib_slim_python_slim_learning.py]

相应的关于如何从一个ckpt文件里导出模型的部分参数用于fine\_tuning:

*************************************************
    * Fine-Tuning Part of a model from a checkpoint *
    *************************************************
    Rather than initializing all of the weights of a given model, we sometimes
    only want to restore some of the weights from a checkpoint. To do this, one
    need only filter those variables to initialize as follows:
      ...
      # Create the train_op
      train_op = slim.learning.create_train_op(total_loss, optimizer)
      checkpoint_path = '/path/to/old_model_checkpoint'
      # Specify the variables to restore via a list of inclusion or exclusion
      # patterns:
      variables_to_restore = slim.get_variables_to_restore(
          include=["conv"], exclude=["fc8", "fc9])
      # or
      variables_to_restore = slim.get_variables_to_restore(exclude=["conv"])
      init_assign_op, init_feed_dict = slim.assign_from_checkpoint(
          checkpoint_path, variables_to_restore)
      # Create an initial assignment function.
      def InitAssignFn(sess):
          sess.run(init_assign_op, init_feed_dict)
      # Run training.
      slim.learning.train(train_op, my_log_dir, init_fn=InitAssignFn)

### Update: ###

使用如下代码：

import tensorflow as tf
    import tensorflow.contrib.slim as slim
    import tensorflow.contrib.slim.nets as nets
    
    images = tf.placeholder(tf.float32, [None, 224, 224, 3])
    predictions = nets.vgg.vgg_16(images)
    print [v.name for v in slim.get_variables_to_restore(exclude=['fc8']) ]

得到输出：

[u'vgg_16/conv1/conv1_1/weights:0',
     u'vgg_16/conv1/conv1_1/biases:0',
     …
     u'vgg_16/fc6/weights:0',
     u'vgg_16/fc6/biases:0',
     u'vgg_16/fc7/weights:0',
     u'vgg_16/fc7/biases:0',
     u'vgg_16/fc8/weights:0',
     u'vgg_16/fc8/biases:0']

如果你不想导出某个层（vgg\_16/fc8）的参数，

print [v.name for v in slim.get_variables_to_restore(exclude=['vgg_16/fc8']) ]

得到输出：

### update2: ###

一个完整的只取部分参数的例子：

import tensorflow as tf
    import tensorflow.contrib.slim as slim
    import tensorflow.contrib.slim.nets as nets
    
    s = tf.Session(config=tf.ConfigProto(gpu_options={'allow_growth':True}))
    
    images = tf.placeholder(tf.float32, [None, 224, 224, 3])
    predictions = nets.vgg.vgg_16(images, 200)
    variables_to_restore = slim.get_variables_to_restore(exclude=['vgg_16/fc8'])
    init_assign_op, init_feed_dict = slim.assign_from_checkpoint('./vgg16.ckpt', variables_to_restore)
    s.run(init_assign_op, init_feed_dict)

以上是从官方提供的slim vgg16在imagenet上训练得到的model中导出部分参数code

[https_github.com_tensorflow_tensorflow_blob_129665119ea60640f7ed921f36db9b5c23455224_tensorflow_contrib_slim_python_slim_learning.py]: https://github.com/tensorflow/tensorflow/blob/129665119ea60640f7ed921f36db9b5c23455224/tensorflow/contrib/slim/python/slim/learning.py