TensorFlow Slim中InceptionV3微调及TFRecord批次处理问题

我最近在尝试用TensorFlow Slim库微调InceptionV3模型，一开始因为官方文档不够完善，写代码时满脑子疑惑，后来硬啃源码才把关键逻辑捋顺，现在已经能顺利完成微调并保存检查点了。下面是我整理的完整操作步骤，希望能帮到有同样困惑的人：

1. 数据准备与读取 #

我已经为训练数据集生成了tf.record格式的文件，目前使用以下代码来读取并预处理数据（完全适配InceptionV3的输入要求）：

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

slim = tf.contrib.slim

# 解析tf.record的自定义函数
def parse_tfrecord(example_proto):
    feature_description = {
        'image/encoded': tf.FixedLenFeature([], tf.string),
        'image/label': tf.FixedLenFeature([], tf.int64),
    }
    features = tf.parse_single_example(example_proto, feature_description)
    
    # 解码图像并执行标准化预处理
    image = tf.image.decode_jpeg(features['image/encoded'], channels=3)
    image = tf.image.resize_images(image, [299, 299])  # InceptionV3要求固定输入尺寸299x299
    image = nets.inception.preprocess_input(image)  # 复用官方预处理逻辑，确保和预训练模型输入一致
    
    # 转换标签为one-hot格式（替换num_classes为你的数据集类别数）
    num_classes = 10
    label = tf.one_hot(features['image/label'], num_classes)
    return image, label

# 构建可迭代的数据集
dataset = tf.data.TFRecordDataset('./train_data.tfrecord')
dataset = dataset.map(parse_tfrecord)
dataset = dataset.shuffle(buffer_size=1000).batch(32).repeat()
iterator = dataset.make_initializable_iterator()
images, labels = iterator.get_next()

2. 加载预训练模型与构建微调网络 #

这一步的核心是合理选择要冻结和微调的网络层，避免破坏预训练的底层特征：

# 构建InceptionV3网络结构
with slim.arg_scope(nets.inception.inception_v3_arg_scope()):
    logits, end_points = nets.inception.inception_v3(
        images,
        num_classes=num_classes,
        is_training=True
    )

# 定义损失函数
loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)

# 选择需要加载预训练权重的变量（排除顶层分类器，因为我们的类别数和预训练模型不同）
restore_vars = slim.get_variables_to_restore(exclude=['InceptionV3/Logits', 'InceptionV3/AuxLogits'])

# 初始化预训练权重加载函数
init_fn = slim.assign_from_checkpoint_fn(
    './inception_v3.ckpt',  # 替换为你的预训练权重文件路径
    restore_vars,
    ignore_missing_vars=True
)

# 选择要微调的变量：这里优先训练顶层分类器和最后一组混合层
train_vars = slim.get_variables_in_scope('InceptionV3/Mixed_7c') + \
             slim.get_variables_in_scope('InceptionV3/Logits') + \
             slim.get_variables_in_scope('InceptionV3/AuxLogits')

3. 配置训练流程与保存检查点 #

最后设置优化器、训练操作，启动会话执行训练并定期保存检查点：

# 配置优化器（用较小的学习率，避免冲掉预训练的有效特征）
optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
train_op = slim.learning.create_train_op(loss, optimizer, variables_to_train=train_vars)

# 创建检查点保存器，保留最近5个版本
saver = tf.train.Saver(max_to_keep=5)

# 启动训练会话
with tf.Session() as sess:
    # 初始化所有变量
    sess.run(tf.global_variables_initializer())
    # 加载预训练权重
    init_fn(sess)
    # 初始化数据集迭代器
    sess.run(iterator.initializer)
    
    # 开始训练循环
    total_steps = 10000
    for step in range(total_steps):
        _, current_loss = sess.run([train_op, loss])
        if step % 100 == 0:
            print(f"Step {step}, 当前损失值: {current_loss:.4f}")
        # 每1000步保存一次检查点
        if step % 1000 == 0 and step != 0:
            saver.save(sess, './fine_tuned_model/inception_v3', global_step=step)
    
    # 训练结束后保存最终模型
    saver.save(sess, './fine_tuned_model/inception_v3_final')
    print("微调完成，模型已成功保存！")

实操小贴士 #

• 预训练权重可以从TensorFlow官方模型仓库获取，注意要匹配InceptionV3的版本
• 微调初期建议只训练顶层分类器，待损失稳定后再逐步放开更多底层网络，能有效降低过拟合风险
• 如果遇到权重加载报错，ignore_missing_vars=True参数可以忽略因类别数不同导致的顶层变量缺失问题
• 图像预处理一定要复用slim.nets.inception.preprocess_input，确保和预训练模型的输入归一化逻辑完全一致

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