TensorFlow训练线性模型：无需TensorBoard绘制迭代-损失曲线

可以通过自定义Keras Callback来记录每个训练批次的损失，之后用Matplotlib绘制损失随迭代次数变化的曲线，具体实现如下：

1. 自定义批次损失记录回调 #

创建一个继承自tf.keras.callbacks.Callback的类，在训练开始时初始化存储损失的列表，每个批次训练结束时记录当前损失：

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

class BatchLossCallback(tf.keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.batch_losses = []
    
    def on_train_batch_end(self, batch, logs=None):
        # 捕获当前批次的训练损失
        self.batch_losses.append(logs['loss'])

2. 修改训练代码，使用自定义回调 #

在model.fit()中传入自定义回调，替换原TensorBoard回调（如果不需要的话）：

d = 1000
N = 10000
X = np.random.rand(N,d) 
theta = np.random.uniform(-2,2,[d,1])
y = X@theta + 2           
y = y.reshape(y.size,) + np.random.randn(y.size)

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1, input_shape=[d,],use_bias=False)])
model.compile(optimizer='adam', loss="mse")

# 初始化自定义回调
batch_loss_callback = BatchLossCallback()
# 执行训练，收集批次损失
history = model.fit(X, y, epochs=1, batch_size=32, callbacks=[batch_loss_callback])

3. 绘制损失曲线 #

利用收集到的批次损失数据，用Matplotlib绘制迭代次数（批次序号）与损失的关系图：

plt.figure(figsize=(10, 6))
plt.plot(range(len(batch_loss_callback.batch_losses)), batch_loss_callback.batch_losses, color='#1f77b4')
plt.xlabel('迭代次数（批次）')
plt.ylabel('训练损失（MSE）')
plt.title('训练损失随迭代次数变化曲线')
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.show()

补充说明 #

• 如果训练多轮（epochs>1），该回调会自动记录所有轮次的批次损失，绘图时X轴为总迭代次数（所有轮次的批次总和）。
• 若需要区分不同轮次的损失，可以在回调中增加epoch记录逻辑，例如在on_epoch_begin中标记当前epoch，再将损失与epoch关联存储。

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