在使用Keras train_on_batch时将直方图摘要添加到tensorboard

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# 创建一个Sequential模型
model = Sequential()
model.add(Dense(64, activation='relu', input_dim=100))
model.add(Dense(64, activation='relu'))
model.add(Dense(10, activation='softmax'))

# 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 创建一个FileWriter对象，用于将摘要写入TensorBoard日志文件
log_dir = 'logs/'
file_writer = tf.summary.create_file_writer(log_dir)

# 定义一个函数，用于将直方图摘要添加到TensorBoard
def add_histogram_summary(step, name, tensor):
    with file_writer.as_default():
        tf.summary.histogram(name, tensor, step=step)
        file_writer.flush()

# 模拟训练过程
for step in range(100):
    # 生成随机输入和标签
    inputs = tf.random.normal((32, 100))
    labels = tf.random.uniform((32, 10))

    # 在每个batch上进行训练
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = model.loss(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    # 将直方图摘要添加到TensorBoard
    add_histogram_summary(step, 'weights', model.weights[0])
    add_histogram_summary(step, 'biases', model.weights[1])

# 关闭FileWriter对象
file_writer.close()