目的：加载权重文件并删除权重文件中最后一层的训练参数
过程描述：由于使用迁移学习的训练方法，加载的官网提供的预训练模型参数，有关网络最后一层的参数并不需要，因此要删掉后才能拿来训练自己的网络
python代码：
import tensorflow as tf
def rename_var(ckpt_path, new_ckpt_path, num_classes, except_list):
with tf.Graph().as_default(), tf.compat.v1.Session().as_default() as sess:
var_list = tf.train.list_variables(ckpt_path)
new_var_list = []
for var_name, shape in var_list:
# print(var_name)
if var_name in except_list:
continue
if "RMSProp" in var_name or "Exponential" in var_name:
continue
var = tf.train.load_variable(ckpt_path, var_name)
new_var_name = var_name.replace('MobilenetV2/', "")
new_var_name = new_var_name.replace("/expand/weights", "/expand/Conv2d/weights")
new_var_name = new_var_name.replace("Conv/weights", "Conv/Conv2d/kernel")
new_var_name = new_var_name.replace("Conv_1/weights", "Conv_1/Conv2d/kernel")
new_var_name = new_var_name.replace("weights", "kernel")
new_var_name = new_var_name.replace("biases", "bias")
first_word = new_var_name.split('/')[0]
if "expanded_conv" in first_word:
last_word = first_word.split('expanded_conv')[-1]
if len(last_word) > 0:
new_word = "inverted_residual" + last_word + "/expanded_conv/"
else:
new_word = "inverted_residual/expanded_conv/"
new_var_name = new_word + new_var_name.split('/', maxsplit=1)[-1]
print(new_var_name)
re_var = tf.Variable(var, name=new_var_name)
new_var_list.append(re_var)
re_var = tf.Variable(tf.keras.initializers.he_uniform()([1280, num_classes]), name="Logits/kernel")
new_var_list.append(re_var)
re_var = tf.Variable(tf.keras.initializers.he_uniform()([num_classes]), name="Logits/bias")
new_var_list.append(re_var)
tf.keras.initializers.he_uniform()
saver = tf.compat.v1.train.Saver(new_var_list)
sess.run(tf.compat.v1.global_variables_initializer())
saver.save(sess, save_path=new_ckpt_path, write_meta_graph=False, write_state=False)
def main():
except_list = ['global_step', 'MobilenetV2/Logits/Conv2d_1c_1x1/biases', 'MobilenetV2/Logits/Conv2d_1c_1x1/weights']
ckpt_path = './pretrain_model/mobilenet_v2_1.0_224.ckpt'
new_ckpt_path = './pretrain_weights.ckpt'
num_classes =4
rename_var(ckpt_path, new_ckpt_path, num_classes, except_list)
if __name__ == '__main__':
main()
代码结果：
?
