TensorFlow：如何按固定时长而非Epochs训练模型？

由于单Epoch耗时波动较大，依赖估算Epoch数的方式不可靠，基于Keras官方的Callback扩展机制，可实现精准的固定时长训练，具体实现如下：

自定义TimeStopping回调类 #

通过继承keras.callbacks.Callback，在训练启动时记录时间戳，每完成一个训练批次后检查已用时，达到设定阈值则终止训练：

import time
from tensorflow import keras

class TimeStopping(keras.callbacks.Callback):
    def __init__(self, total_seconds):
        super().__init__()
        self.total_seconds = total_seconds
        self.start_time = None

    def on_train_begin(self, logs=None):
        self.start_time = time.time()

    def on_batch_end(self, batch, logs=None):
        elapsed_time = time.time() - self.start_time
        if elapsed_time >= self.total_seconds:
            self.model.stop_training = True
            print(f"\n已达到训练时长上限{self.total_seconds}秒，终止训练")

集成到现有训练流程 #

将自定义回调加入训练的callbacks列表，例如设定训练时长为3600秒（1小时）：

# 原有代码保持不变
y = to_categorical(self.ydata, num_classes=self.vocab_size)
model = Sequential()
model.add(Embedding(self.vocab_size, 10, input_length=1))
model.add(LSTM(1000, return_sequences=True))
model.add(LSTM(1000))
model.add(Dense(1000, activation="relu"))
model.add(Dense(self.vocab_size, activation="softmax"))
keras.utils.plot_model(model, show_layer_names=True)
checkpoint = ModelCheckpoint(modelFilePath, monitor='loss', verbose=1,save_best_only=True, mode='auto')
reduce = ReduceLROnPlateau(monitor='loss', factor=0.2,patience=3, min_lr=0.0001, verbose=1)
tensorboard_Visualization = TensorBoard(log_dir=logdirPath)

# 初始化固定时长回调
time_stopper = TimeStopping(total_seconds=3600)

model.compile(loss="categorical_crossentropy", optimizer=Adam(lr=0.001))
# 将time_stopper加入回调列表，epochs设置为足够大的值确保训练被时间回调终止
history = model.fit(
    self.Xdata, y, 
    epochs=1000,
    batch_size=64, 
    callbacks=[checkpoint, reduce, tensorboard_Visualization, time_stopper]
).history

方案说明 #

• 该方案基于Keras官方回调机制，属于官方支持的训练流程扩展方式，稳定性与兼容性有保障。
• 选择在on_batch_end节点检查时长，相比on_epoch_end能更精准控制训练时长，避免单个Epoch耗时过长导致超时。
• 训练时需将epochs设置为足够大的数值，确保训练不会因Epoch数耗尽而提前终止。

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