如何在TensorBoard中可视化各训练Epoch的MxN张量输出

TensorBoard支持这种需求，你可以通过自定义Keras回调结合tf.summary写入序列数据，实现按Epoch查看N条时间序列曲线的效果，具体操作如下：

1. 编写自定义回调处理张量日志 #

创建回调类，在每个Epoch结束时获取模型输出的MxN张量，将其拆分为N个长度为M的序列，分别写入TensorBoard日志：

import tensorflow as tf
from tensorflow.keras import callbacks

class TensorSequenceLogger(callbacks.Callback):
    def __init__(self, log_dir, test_input):
        super().__init__()
        self.writer = tf.summary.create_file_writer(log_dir)
        self.test_input = test_input  # 用于获取模型输出的测试数据
    
    def on_epoch_end(self, epoch, logs=None):
        # 获取当前Epoch的MxN张量输出
        tensor_out = self.model.predict(self.test_input, verbose=0)  # shape (M, N)
        # 转置为(N, M)，方便遍历每个序列
        seq_list = tf.transpose(tensor_out)
        
        with self.writer.as_default():
            for seq_id, seq_data in enumerate(seq_list):
                tag = f"tensor_sequences/seq_{seq_id}"
                # 为每个序列的每个时间步写入标量，用epoch*M + step作为全局step，避免不同Epoch重叠
                for step, val in enumerate(seq_data):
                    global_step = epoch * tensor_out.shape[0] + step
                    tf.summary.scalar(tag, val, step=global_step)

2. 训练时挂载回调 #

在模型训练时同时添加原生TensorBoard回调和自定义回调：

log_dir = "./tensorboard_logs"
test_input = ...  # 替换为你的测试输入数据

# 原生TensorBoard回调
tb_callback = callbacks.TensorBoard(log_dir=log_dir)
# 自定义张量序列日志回调
seq_callback = TensorSequenceLogger(log_dir=log_dir, test_input=test_input)

model.fit(
    train_dataset,
    epochs=30,
    callbacks=[tb_callback, seq_callback]
)

3. 在TensorBoard中查看 #

启动TensorBoard后进入SCALARS标签页，找到tensor_sequences分组下的N条曲线。你可以通过以下方式查看指定Epoch的序列：

• 使用图例的筛选功能，勾选需要查看的序列
• 调整X轴范围，定位到目标Epoch对应的区间（每个Epoch占M个step长度）

如果希望更便捷地切换Epoch，也可以将每个Epoch的张量用tf.summary.tensor写入日志，在TensorBoard的TENSORS标签页查看完整张量数据，再导出后用Matplotlib等工具绘制曲线，但这种方式无法直接在TensorBoard内交互查看曲线。

注意：global_step = epoch * tensor_out.shape[0] + step中的tensor_out.shape[0]就是M，确保每个Epoch的step区间不重叠即可。

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