Tensorflow 预训练权重有选择加载参数

训练自己的模型我们一般都会使用预训练权重，有时别人训练好的模型类别与自己的不符合，这个时候就不需要加载一些层的参数.

      由于原网络权重为150类，我只需要几类，最后一层的conv带有num_class参数，所以“your name scope” 替换为不需要的层的name,加载模型的时候就不会加载这层的参数。

sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

...........................................

restore_var = [v for v in tf.global_variables()]
variables_to_resotre = [v for v in restore_var if v.name.split('/')[0] != 'your name scope']

..............................................

ckpt = tf.train.get_checkpoint_state('./model/')
saver = tf.train.Saver(var_list=variables_to_resotre)
saver.restore(sess, ckpt.model_checkpoint_path)

下面给一段查看各个网络层参数的name scope, 便于修改。

from tensorflow.python import pywrap_tensorflow


checkpoint_path = “your ckpt path”
# Read data from checkpoint file

reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)

var_to_shape_map = reader.get_variable_to_shape_map()

# Print tensor name and values

for key in var_to_shape_map:

    print("tensor_name: ", key)

    print(reader.get_tensor(key))

 还可以使用tensorflow内置模块slim加载需要的层

exclude = ['layer1', 'layer2']
variables_to_restore = slim.get_variables_to_restore(exclude=exclude)
saver = tf.train.Saver(variables_to_restore)
saver.restore(sess, “model_path”)