有人设法用 TensorBoard 记录损失吗?我正在使用 trax ml 库。我收到此错误TypeError: 'SummaryWriter' object is not callable
。
我正在使用SummaryWriter
from jaxboard
,然后将其添加到callbacks
inside training.Loop
。
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())