如何为TensorFlow模型.fit()日志添加时间戳及获取epoch训练时长

问题描述 #

使用TensorFlow 2.15训练模型时，.fit()输出的日志示例如下：

Epoch 2/10
32/32 - 0s - loss: 17.3287 - 99ms/epoch - 3ms/step
Epoch 3/10
32/32 - 0s - loss: 16.9345 - 123ms/epoch - 4ms/step

需求是为这类日志添加时间戳，但修改Python logging模块的格式后，.fit()的输出没有变化。尝试用on_epoch_end回调自定义日志时，无法获取ms/step这类运行时指标，现咨询：

1. 能否直接为.fit()的输出日志添加时间戳？
2. 若不可行，如何在on_epoch_end回调中获取每个epoch的训练时长？

附带的测试代码如下：

import logging
import sys

import keras
import numpy as np
import tensorflow as tf

# logging setup
logger = tf.get_logger()
handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)


class LoggerTensorflow(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        epoch_string = f"Epoch {epoch + 1} / {self.params['epochs']}"
        logs_formatted = " - ".join([f"{k}: {v:.4f}" for k, v in logs.items()])
        logger.info(f"{epoch_string} - {logs_formatted}")
        return super().on_epoch_end(epoch, logs)


# Generate some fake data: 1000 samples with 1 feature.
np.random.seed(42)  # for reproducible results
x_data = np.random.rand(1000, 1)  # features
y_data = 3 * x_data + 2 + np.random.normal(0, 0.05, (1000, 1))  # targets with noise


# Build a simple linear regression model
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

# Compile the model specifying the optimizer, loss function, and metrics to track
model.compile(loss="mse")

# Train the model
model.fit(x_data, y_data, epochs=10, verbose=2)

解答 #

1. 直接修改.fit()默认日志格式添加时间戳？不行 #

TensorFlow的.fit()默认输出（由verbose参数控制）是直接打印到标准输出的，并不经过你配置的Python logging模块，所以修改logging.Formatter不会影响这部分内容。无法直接修改.fit()默认日志的格式来添加时间戳，必须通过自定义回调替代默认的日志输出逻辑。

2. 在on_epoch_end回调中获取epoch训练时长的实现方法 #

通过在on_epoch_begin记录当前时间，在on_epoch_end计算时间差，就能得到该epoch的训练时长。同时，ms/step可以通过总时长除以该epoch的步数（self.params['steps']）计算得出。

修改后的自定义回调及完整测试代码如下：

import logging
import sys
import time

import keras
import numpy as np
import tensorflow as tf

# logging setup
logger = tf.get_logger()
handler = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)  # 确保INFO级别日志能输出


class LoggerTensorflow(keras.callbacks.Callback):
    def on_epoch_begin(self, epoch, logs=None):
        # 记录epoch开始时间
        self.start_time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        epoch_num = epoch + 1
        total_epochs = self.params['epochs']
        steps = self.params['steps']
        
        # 计算时长（转换为毫秒）
        epoch_duration = (time.time() - self.start_time) * 1000
        ms_per_step = epoch_duration / steps

        # 格式化日志内容
        epoch_string = f"Epoch {epoch_num}/{total_epochs}"
        progress_string = f"{steps}/{steps} - {epoch_duration/1000:.1f}s"
        logs_formatted = " - ".join([f"{k}: {v:.4f}" for k, v in logs.items()])
        timing_string = f"- {epoch_duration:.0f}ms/epoch - {ms_per_step:.0f}ms/step"
        
        # 组合所有内容并输出
        full_log = f"{epoch_string}\n{progress_string} - {logs_formatted} {timing_string}"
        logger.info(full_log)


# Generate some fake data: 1000 samples with 1 feature.
np.random.seed(42)  # for reproducible results
x_data = np.random.rand(1000, 1)  # features
y_data = 3 * x_data + 2 + np.random.normal(0, 0.05, (1000, 1))  # targets with noise


# Build a simple linear regression model
model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

# Compile the model specifying the optimizer, loss function, and metrics to track
model.compile(loss="mse")

# 训练时设置verbose=0，关闭默认输出，只使用自定义回调的日志
model.fit(x_data, y_data, epochs=10, verbose=0, callbacks=[LoggerTensorflow()])

代码说明 #

• 新增on_epoch_begin方法记录epoch开始的时间戳
• 在on_epoch_end中计算epoch总时长（毫秒）和单步耗时（ms/step）
• 设置verbose=0关闭.fit()的默认输出，避免重复日志
• 自定义日志格式完全匹配原.fit()的输出结构，并添加了时间戳

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