如何计算TensorFlow中模型全训练周期（所有Epoch）的总耗时？

自定义回调记录每个Epoch耗时并累加 #

通过继承tf.keras.callbacks.Callback实现自定义回调，在每个Epoch的开始和结束节点记录时间，自动累加总耗时：

import tensorflow as tf
import time

class TimeHistory(tf.keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.total_training_time = 0.0

    def on_epoch_begin(self, epoch, logs=None):
        self._epoch_start = time.time()

    def on_epoch_end(self, epoch, logs=None):
        epoch_duration = time.time() - self._epoch_start
        self.total_training_time += epoch_duration
        print(f"Epoch {epoch+1} 耗时: {epoch_duration:.2f}s")

# 初始化回调实例
time_cb = TimeHistory()

# 训练时传入回调
model.fit(
    train_dataset,
    epochs=YOUR_EPOCH_NUM,
    callbacks=[time_cb]
)

# 训练结束后输出总耗时
print(f"全量数据集训练总耗时: {time_cb.total_training_time:.2f}s")

简化方案：直接记录训练全程时间 #

如果不需要单独统计每个Epoch的耗时，直接在训练前后记录时间差即可：

import time

# 训练前记录起始时间
start_total = time.time()

model.fit(
    train_dataset,
    epochs=YOUR_EPOCH_NUM
)

# 计算并输出总耗时
total_duration = time.time() - start_total
print(f"全量数据集训练总耗时: {total_duration:.2f}s")

注意点 #

• 以上方法包含数据加载、预处理和模型训练的全部时间；若需仅统计模型正向/反向传播时间，需在数据 pipeline 和训练环节拆分计时
• 分布式训练场景下，time.time()仅记录单进程时间，建议改用tf.timestamp()来适配分布式环境的时间统计

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