如何在TensorFlow 2的model_main_tf2.py中实现早停回调？

针对使用model_main_tf2.py训练时无法直接用model.fit()添加EarlyStopping的问题，以下是两种可行的解决方案：

方案一：修改model_main_tf2.py集成Keras原生早停回调 #

1. 导入必要模块
   在model_main_tf2.py文件开头添加：

from tensorflow.keras.callbacks import EarlyStopping

2. 配置早停参数
   在main函数中，调用默认训练循环前，定义早停规则：

# 监控验证集损失，连续3个epoch无下降则停止训练，恢复最优权重
early_stopping = EarlyStopping(
    monitor='val_loss',
    patience=3,
    restore_best_weights=True,
    verbose=1
)

3. 替换默认训练循环，手动实现带回调的流程
   默认的model_lib_v2.train_loop未集成回调逻辑，需替换为自定义训练循环，手动触发回调的生命周期方法：

# 加载配置文件
pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
with tf.io.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f:
  text_format.Merge(f.read(), pipeline_config)
train_config = pipeline_config.train_config
eval_config = pipeline_config.eval_config

# 计算单epoch步数（训练集总数//batch_size，需根据你的数据调整）
steps_per_epoch = 1000  # 示例值，替换为实际计算值

# 构建训练、验证输入函数和模型
train_input_fn = model_lib_v2.create_train_input_fn(pipeline_config)
eval_input_fn = model_lib_v2.create_eval_input_fn(pipeline_config)
model = model_lib_v2.build_model(pipeline_config.model, is_training=True)
optimizer = model_lib_v2.create_optimizer(train_config)

# 初始化早停回调
early_stopping.on_train_begin()

# 自定义训练循环
for epoch in range(train_config.num_epochs or FLAGS.num_train_steps // steps_per_epoch):
    # 训练一个epoch
    train_loss = model_lib_v2.train_one_epoch(
        model, optimizer, train_input_fn, steps_per_epoch
    )
    
    # 运行验证
    eval_results = model_lib_v2.evaluate(
        model, eval_input_fn, eval_config.num_examples
    )
    val_loss = eval_results['loss']
    
    # 触发epoch结束回调
    logs = {'loss': train_loss, 'val_loss': val_loss}
    early_stopping.on_epoch_end(epoch, logs)
    
    # 检查是否触发早停
    if early_stopping.model.stop_training:
        print("早停触发，终止训练")
        break
    
    # 定期保存模型
    if epoch % FLAGS.checkpoint_every_n == 0:
        model.save_weights(f"{FLAGS.model_dir}/epoch_{epoch}")

方案二：自定义早停逻辑（最小化修改原文件） #

如果不想大幅改动原代码，可以自己实现轻量早停类，在训练中手动检查指标：

1. 定义自定义早停类
   在model_main_tf2.py中添加：

class CustomEarlyStopping:
    def __init__(self, monitor='val_loss', patience=3, restore_best_weights=True):
        self.monitor = monitor
        self.patience = patience
        self.restore_best_weights = restore_best_weights
        # 根据指标类型初始化最优值（loss越小越好，acc/mAP越大越好）
        self.best_score = float('inf') if 'loss' in monitor else -float('inf')
        self.wait_count = 0
        self.best_model_weights = None

    def check_stop(self, model, current_score):
        # 判断当前指标是否更优
        if ('loss' in self.monitor and current_score < self.best_score) or \
           ('map' in self.monitor and current_score > self.best_score):
            self.best_score = current_score
            self.wait_count = 0
            if self.restore_best_weights:
                self.best_model_weights = model.get_weights()
        else:
            self.wait_count += 1
            if self.wait_count >= self.patience:
                print(f"连续{self.patience}个epoch无指标提升，触发早停")
                if self.restore_best_weights:
                    model.set_weights(self.best_model_weights)
                return True
        return False

2. 集成到训练流程
   替换原train_loop调用，加入早停检查：

# 加载配置、模型、输入函数（同方案一）
early_stopper = CustomEarlyStopping(patience=5, monitor='val_map')

for epoch in range(...):
    train_loss = model_lib_v2.train_one_epoch(...)
    eval_results = model_lib_v2.evaluate(...)
    current_val_map = eval_results['DetectionBoxes_Precision/mAP']
    
    # 检查是否触发早停
    if early_stopper.check_stop(model, current_val_map):
        break

注意事项 #

• 确保你的pipeline配置文件中已正确配置验证集路径和评价参数，否则无法获取验证指标触发早停。
• 可根据需求切换监控指标，比如目标检测常用的val_map，对应修改比较逻辑即可。

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