如何基于步数间隔而非时间间隔保存TensorFlow-Slim模型？

不用修改TensorFlow-Slim的源码就能实现按步数间隔保存模型，我之前做目标检测训练时也遇到过这个问题，下面是几个可靠的方案：

方案一：自定义SessionRunHook实现步数触发保存 #

TensorFlow的tf.train.SessionRunHook可以让我们介入训练的各个阶段，正好可以用来监控全局步数，达到指定间隔时触发模型保存。

具体实现步骤：

• 定义一个继承自SessionRunHook的钩子类，在初始化时传入tf.train.Saver、保存间隔步数和保存目录
• 在begin()方法中获取全局步数的张量，确保训练图中有这个张量（一般训练时都会创建）
• 在after_run()方法中，获取当前训练步数，判断是否达到保存间隔，如果是就调用saver的保存方法

代码示例：

import tensorflow as tf
from tensorflow.contrib import slim

class StepBasedSaverHook(tf.train.SessionRunHook):
    def __init__(self, saver, save_steps, save_dir):
        self.saver = saver
        self.save_steps = save_steps
        self.save_dir = save_dir
        self.last_saved_step = 0

    def begin(self):
        # 获取全局步数张量
        self.global_step = tf.train.get_global_step()
        if self.global_step is None:
            raise ValueError("训练图中未创建全局步数张量，请确保定义了tf.train.get_or_create_global_step()")

    def after_run(self, run_context, run_values):
        # 获取当前步数
        current_step = run_context.session.run(self.global_step)
        # 判断是否达到保存间隔
        if current_step - self.last_saved_step >= self.save_steps:
            self.saver.save(
                run_context.session,
                self.save_dir,
                global_step=current_step
            )
            self.last_saved_step = current_step
            print(f"已保存模型到步数 {current_step}")

# --- 训练代码部分 ---
# 假设你已经定义了train_op和模型结构
log_dir = "./training_checkpoints"
save_every_steps = 1000  # 每1000步保存一次

# 创建Saver实例
saver = tf.train.Saver(max_to_keep=5)
# 初始化自定义钩子
step_saver_hook = StepBasedSaverHook(saver, save_every_steps, log_dir)

# 启动训练，传入自定义钩子，并将默认的时间间隔保存设为很大的值避免冲突
slim.learning.train(
    train_op=train_op,
    logdir=log_dir,
    saver=saver,
    hooks=[step_saver_hook],
    save_interval_secs=999999  # 禁用默认的时间间隔保存
)

方案二：使用CheckpointManager配合Hook（更现代的方式） #

如果你偏向使用TensorFlow更现代的检查点管理方式，可以用tf.train.CheckpointManager来替代传统的tf.train.Saver，同样通过Hook触发保存：

class StepBasedCheckpointHook(tf.train.SessionRunHook):
    def __init__(self, checkpoint, save_steps, save_dir):
        self.checkpoint = checkpoint
        self.save_steps = save_steps
        self.manager = tf.train.CheckpointManager(checkpoint, save_dir, max_to_keep=5)
        self.last_saved_step = 0

    def begin(self):
        self.global_step = tf.train.get_global_step()

    def after_run(self, run_context, run_values):
        current_step = run_context.session.run(self.global_step)
        if current_step - self.last_saved_step >= self.save_steps:
            self.manager.save(checkpoint_number=current_step)
            self.last_saved_step = current_step
            print(f"已保存检查点到步数 {current_step}")

# --- 训练代码部分 ---
# 创建Checkpoint实例，包含模型、优化器和全局步数
checkpoint = tf.train.Checkpoint(
    model=your_model,
    optimizer=your_optimizer,
    global_step=tf.train.get_global_step()
)
checkpoint_hook = StepBasedCheckpointHook(checkpoint, 1000, "./training_checkpoints")

slim.learning.train(
    train_op=train_op,
    logdir=log_dir,
    hooks=[checkpoint_hook],
    save_interval_secs=999999,
    saver=None  # 不需要传入传统Saver了
)

不推荐的临时方案：估算时间间隔 #

如果你只是临时测试，可以根据你的训练速度估算每n步需要的时间，然后设置save_interval_secs参数。比如你每步耗时0.1秒，1000步就是100秒，就设置save_interval_secs=100。但这个方法很不可靠，一旦训练速度波动（比如数据加载变慢、GPU负载变化），保存的步数就会不准，只适合短期测试用。

为什么不用修改Slim源码？ #

因为Slim的slim.learning.train()函数本身就支持传入hooks参数，而SessionRunHook是TensorFlow原生的扩展机制，允许我们在不修改框架源码的情况下，介入训练的生命周期，实现自定义的逻辑（包括按步数保存）。

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