使用Trax的TensorBoard

import trax
from trax.supervised import training
import tensorflow as tf

# 定义模型
model = trax.models.Sequential([
    trax.layers.Dense(512, activation=trax.layers.Relu),
    trax.layers.Dense(10, activation=trax.layers.Softmax)
])

# 定义损失函数和优化器
loss_fn = trax.losses.SparseCategoricalCrossentropy()
optimizer = trax.optimizers.Adam(learning_rate=0.001)

# 创建训练任务
train_task = training.TrainTask(
    labeled_data=train_data,
    loss_layer=loss_fn,
    optimizer=optimizer,
    n_steps_per_checkpoint=10
)

# 创建评估任务
eval_task = training.EvalTask(
    labeled_data=eval_data,
    metrics=[trax.metrics.Accuracy()]
)

# 创建训练循环
training_loop = training.Loop(
    model,
    train_task,
    eval_tasks=[eval_task],
    output_dir='./output'
)

# 集成 TensorBoard
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs')

# 启动训练
training_loop.run(n_steps=100, callbacks=[tensorboard_callback])