TensorFlow中预训练VGG19特定权重加载、特征提取与微调问题

针对你提到的两个核心问题，我会结合TensorFlow的实践给出具体的实现方案和代码示例：

一、从VGG19 ckpt中提取特定层的权重 #

要提取特定层的权重，首先得明确ckpt文件中包含的变量名，然后筛选出你需要的部分：

1. 查看ckpt中的所有变量名
   先用以下代码列出ckpt里的所有变量，方便你定位目标层：

   import tensorflow as tf
   
   ckpt_path = "path/to/vgg_19.ckpt"
   var_list = tf.train.list_variables(ckpt_path)
   for var_name, shape in var_list:
       print(var_name, shape)
   

   比如VGG19的卷积层变量名通常是vgg_19/conv1/conv1_1/weights、vgg_19/conv1/conv1_1/biases这种格式，全连接层则是vgg_19/fc6/weights等。

2. 筛选并提取目标层权重
   如果你只想保留所有卷积层的权重（可根据需求调整过滤规则），可以构建一个变量名过滤列表，然后用tf.train.Saver加载并保存这些变量：

   # 定义需要保留的变量名前缀，比如所有卷积层和池化层
   target_prefixes = ["vgg_19/conv", "vgg_19/pool"]
   
   # 获取符合条件的变量列表
   target_vars = [var for var in tf.global_variables() if any(var.name.startswith(prefix) for prefix in target_prefixes)]
   
   # 加载并保存提取的权重到新的ckpt
   saver = tf.train.Saver(target_vars)
   with tf.Session() as sess:
       saver.restore(sess, ckpt_path)
       saver.save(sess, "path/to/extracted_vgg_conv.ckpt")
   

二、分阶段加载权重并进行微调 #

这里分两种场景：首次训练（加载预训练部分权重，其余随机初始化）、后续训练（加载全量训练后的权重）

场景1：首次训练 - 加载预训练特定层，其余随机初始化 #

首先构建你的自定义模型（基于VGG19卷积层+新输出层，或微调最后一层卷积），然后指定要从预训练ckpt恢复的变量：

import tensorflow as tf
from tensorflow.contrib.slim.nets import vgg

# 构建模型
images = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
with tf.contrib.slim.arg_scope(vgg.vgg_arg_scope()):
    # 假设你要微调最后一层卷积（conv5_4），前面的卷积层固定
    logits, end_points = vgg.vgg_19(images, num_classes=1000, is_training=True)

# 筛选需要从预训练ckpt恢复的变量：排除最后一层卷积和全连接层
restore_vars = []
for var in tf.global_variables():
    if not var.name.startswith("vgg_19/conv5/conv5_4") and not var.name.startswith("vgg_19/fc"):
        restore_vars.append(var)

# 定义恢复用的saver（只恢复指定变量）和全量保存用的saver（用于后续训练）
restore_saver = tf.train.Saver(restore_vars)
full_saver = tf.train.Saver()

with tf.Session() as sess:
    # 先初始化所有变量（未被预训练覆盖的变量会随机初始化）
    sess.run(tf.global_variables_initializer())
    
    # 从预训练ckpt恢复指定变量
    restore_saver.restore(sess, "path/to/vgg_19.ckpt")
    
    # 定义训练操作：只训练目标变量（最后一层卷积和全连接层）
    train_vars = [var for var in tf.global_variables() if var.name.startswith("vgg_19/conv5/conv5_4") or var.name.startswith("vgg_19/fc")]
    loss = ...  # 根据你的任务定义损失函数
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
    train_op = optimizer.minimize(loss, var_list=train_vars)
    
    # 训练循环（省略具体迭代逻辑）
    # ...
    
    # 训练结束后保存全量权重
    full_saver.save(sess, "path/to/trained_model.ckpt")

场景2：后续训练 - 加载全量训练后的权重 #

后续训练时，直接用全量的Saver加载即可，不需要再区分变量：

import tensorflow as tf
from tensorflow.contrib.slim.nets import vgg

# 重新构建与首次训练完全一致的模型
images = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
with tf.contrib.slim.arg_scope(vgg.vgg_arg_scope()):
    logits, end_points = vgg.vgg_19(images, num_classes=1000, is_training=True)

# 定义全量saver
full_saver = tf.train.Saver()

with tf.Session() as sess:
    # 直接加载全量训练后的权重
    full_saver.restore(sess, "path/to/trained_model.ckpt")
    
    # 继续训练（可调整学习率等参数）
    # ...

关键注意点 #

• 变量名匹配：确保你构建的模型中变量名和预训练ckpt中的变量名完全一致，否则无法恢复权重。如果自定义了模型结构，可能需要手动调整变量名或用var_list做映射。
• 微调策略：如果要固定前面的卷积层，一定要在优化器中指定只训练目标变量，避免前面的层被意外修改。
• 初始化顺序：首次训练时，必须先运行tf.global_variables_initializer()，再恢复预训练权重，这样未被恢复的变量会保持随机初始化状态。

内容的提问来源于stack exchange，提问作者JohnVS