使用Keras model.fit时，如何在TensorBoard查看各层单独的参数直方图

我有一个包含两个隐藏层的Keras Sequential模型，希望在TensorBoard中分别查看每层的kernel和bias的直方图（即Layer_1/Bias、Layer_1/Kernel、Layer_2/Bias、Layer_2/Kernel各一个独立直方图）。但目前TensorBoard中仅显示一个合并的Kernel直方图和一个合并的Bias直方图，看起来是将两层的参数合并在了一起。

模型定义如下：

model = keras.Sequential()
model.add(keras.Input(shape=(1,)))
model.add(layers.Dense(256, activation='relu', name='layer_1'))
model.add(layers.Dense(256, activation='relu', name='layer_2'))
model.add(layers.Dense(1))
...
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
...
model.fit(xTrain, yTrain, epochs = 50, callbacks=[tensorboard_callback])

使用版本：TensorBoard 2.18.0、Keras 3.7.0、Python 3.11。希望找到不用自定义训练循环，仅用Keras高级model.fit函数的解决方案。

默认TensorBoard回调会将同类型参数（如所有Dense层的kernel）合并展示，要实现分层查看，可通过自定义回调手动记录每层参数的直方图，无需修改训练循环：

1. 定义自定义回调类，遍历模型每层的可训练参数，按「层名/参数名」的格式记录直方图：

import tensorflow as tf
from keras import callbacks

class LayerwiseHistogramCallback(callbacks.Callback):
    def __init__(self, log_dir):
        super().__init__()
        self.log_dir = log_dir
        self.summary_writer = tf.summary.create_file_writer(log_dir)

    def on_epoch_end(self, epoch, logs=None):
        with self.summary_writer.as_default():
            # 遍历模型所有层
            for layer in self.model.layers:
                # 记录kernel直方图（仅当层有可训练kernel时）
                if hasattr(layer, 'kernel') and layer.kernel.trainable:
                    tf.summary.histogram(f'{layer.name}/Kernel', layer.kernel, step=epoch)
                # 记录bias直方图（仅当层有可训练bias时）
                if hasattr(layer, 'bias') and layer.bias.trainable:
                    tf.summary.histogram(f'{layer.name}/Bias', layer.bias, step=epoch)
        self.summary_writer.flush()

2. 修改训练代码，关闭默认TensorBoard的直方图记录（避免重复），同时添加自定义回调：

log_dir = "./logs"
# 原TensorBoard回调保留用于记录其他指标，关闭默认直方图
tensorboard_callback = callbacks.TensorBoard(log_dir=log_dir, histogram_freq=0)
# 实例化自定义分层直方图回调
layer_hist_callback = LayerwiseHistogramCallback(log_dir=log_dir)

# 传入两个回调进行训练
model.fit(xTrain, yTrain, epochs=50, callbacks=[tensorboard_callback, layer_hist_callback])

3. 启动TensorBoard查看结果：

tensorboard --logdir=./logs

在TensorBoard的「Histograms」标签下，即可看到layer_1/Kernel、layer_1/Bias、layer_2/Kernel、layer_2/Bias各自独立的直方图。

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