带有 Trax 的 TensorBoard答案

【问题标题】：TensorBoard with Trax带有 Trax 的 TensorBoard
【发布时间】：2021-05-27 13:55:10
【问题描述】：

有人设法用 TensorBoard 记录损失吗？我正在使用 trax ml 库。我收到此错误TypeError: 'SummaryWriter' object is not callable。

我正在使用来自jaxboard 的SummaryWriter，然后将其添加到training.Loop 内的callbacks。

my_dir = "/some_dir" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
summary_writer = jaxboard.SummaryWriter(log_dir=my_dir)

def train_model(model, batch_size=batch_size, n_steps=1000, output_dir=output_dir):
    '''
    Input: 
        model - the model
        train_task - Training task
        eval_tasks - Evaluation task
        n_steps - the evaluation steps
        output_dir - folder to save your files
    Output:
        trainer -  trax trainer
    '''
    train_task = training.TrainTask(
        labeled_data=train_generator(batch_size=batch_size, shuffle=True),
        loss_layer=TripletLoss(),
        optimizer=trax.optimizers.Adam(learning_rate=0.001),
        n_steps_per_checkpoint=1000,
    )

    eval_tasks = training.EvalTask(
        labeled_data=val_generator(batch_size=batch_size, shuffle=True),
        metrics=[TripletLoss()],
        n_eval_batches=10,
    )

    training_loop = training.Loop(
                                model, # The learning model
                                train_task, # The training task
                                eval_tasks = eval_tasks, # The evaluation task
                                #random_seed=35,
                                output_dir = output_dir, # The output directory
                                callbacks=[summary_writer], # Logging
                                ) 
    
    training_loop.run(n_steps = n_steps)
    
    # Return the training_loop, since it has the model.
    return training_loop

运行训练循环时出现错误：

training_loop = train_model(my_model())

【问题讨论】：

标签： tensorflow tensorboard trax

【解决方案1】：

当我删除带有回调的行时起作用，summary_writer 而是在 google colab 上添加了这个：

%load_ext tensorboard
%cd '/content/drive/MyDrive/path_to_the_notebook_/'
# *train* folder is created by trax and holds the logs for the train run
# to log the eval run change to --logdir eval
%tensorboard --logdir train # train = name of folder

【讨论】：