如何阻止TensorFlow Estimator恢复权重与全局训练步数？

这个问题我之前做项目的时候也碰到过，TensorFlow的Estimator默认会把训练的检查点（包括模型权重和global_step）存在你指定的model_dir目录里，每次启动训练时都会自动加载最新的检查点，所以才会出现步数累加、权重自动恢复的情况。要实现每次训练都从零开始初始化，有几个靠谱的解决办法：

方法1：每次训练前清空检查点目录 #

这是最直接的方式，因为Estimator找不到旧的检查点文件，自然就会从头开始初始化所有变量。你可以用Python的shutil模块来删除整个目录，记得处理目录不存在的情况避免报错：

import shutil
import os

model_dir = "./my_model_checkpoint"

# 训练前先清空旧的检查点目录
if os.path.exists(model_dir):
    shutil.rmtree(model_dir)

# 正常创建Estimator并启动训练
estimator = tf.estimator.Estimator(model_fn=your_model_fn, model_dir=model_dir)
estimator.train(input_fn=train_input_fn, steps=2000)

方法2：每次训练使用全新的model_dir #

如果不想删除旧目录（比如需要保留之前的训练记录用于对比），可以每次给model_dir加个唯一标识，比如时间戳或者随机数，这样每个训练任务都会使用独立的存储环境，完全不会受之前训练的影响：

import time

# 用当前时间戳生成唯一的检查点目录名
model_dir = f"./my_model_{int(time.time())}"

estimator = tf.estimator.Estimator(model_fn=your_model_fn, model_dir=model_dir)
estimator.train(input_fn=train_input_fn, steps=2000)

方法3：在model_fn中强制初始化变量（进阶方案） #

如果因为某些限制不能修改model_dir，可以在model_fn里手动控制变量初始化逻辑，跳过检查点的恢复流程。具体来说，当模式为TRAIN时，强制运行全局变量初始化操作，同时让Estimator明确忽略检查点恢复：

def your_model_fn(features, labels, mode):
    # 定义你的模型结构（示例）
    dense_layer = tf.layers.Dense(units=10)
    logits = dense_layer(features)
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
        global_step = tf.train.get_global_step()
        train_op = optimizer.minimize(loss, global_step=global_step)
        
        # 强制初始化所有全局变量，覆盖检查点恢复逻辑
        init_op = tf.global_variables_initializer()
        with tf.control_dependencies([init_op]):
            train_op = tf.identity(train_op)
        
        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            train_op=train_op,
            # 明确设置不使用warm start，避免恢复检查点
            warm_start_settings=tf.estimator.WarmStartSettings(None)
        )
    
    # 处理EVAL和PREDICT模式的逻辑...
    elif mode == tf.estimator.ModeKeys.EVAL:
        eval_metric_ops = {
            "accuracy": tf.metrics.accuracy(labels=labels, predictions=tf.argmax(logits, 1))
        }
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
    elif mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {"class_ids": tf.argmax(logits, 1)}
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

注意：方法3需要确保你的model_fn里没有其他依赖检查点的逻辑，否则可能会出现冲突。一般来说，优先选择前两种方法会更稳妥、更易维护。

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