具有自定义训练循环的 Tensorboard Graph 不包括我的模型

Question

我创建了自己的循环，如 TF 2 迁移指南 here 中所示。
我目前只能看到下面代码的--- VISIBLE --- 部分的图表。如何使我的模型（在 ---NOT VISIBLE--- 部分中定义）在张量板中可见？

如果我没有使用自定义训练循环，我可以使用 documented model.fit approach：

model.fit(..., callbacks=[keras.callbacks.TensorBoard(log_dir=logdir)])

在 TF 1 中，该方法过去非常简单：

tf.compat.v1.summary.FileWriter(LOGDIR, sess.graph)

Tensorboard 迁移指南明确指出 (here)：

不直接编写 tf.compat.v1.Graph - 而是使用 @tf.function 和跟踪函数

configure_default_gpus()
tf.summary.trace_on(graph=True)
K = tf.keras
dataset = sanity_dataset(BATCH_SIZE)

#-------------------------- NOT VISIBLE -----------------------------------------
model = K.models.Sequential([
    K.layers.Flatten(input_shape=(IMG_WIDTH, IMG_HEIGHT, IMG_CHANNELS)),
    K.layers.Dense(10, activation=K.layers.LeakyReLU()),
    K.layers.Dense(IMG_WIDTH * IMG_HEIGHT * IMG_CHANNELS, activation=K.layers.LeakyReLU()),
    K.layers.Reshape((IMG_WIDTH, IMG_HEIGHT, IMG_CHANNELS)),
])
#--------------------------------------------------------------------------------

optimizer = tf.keras.optimizers.Adam()
loss_fn = K.losses.Huber()


@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs, training=True)
#-------------------------- VISIBLE ---------------------------------------------
        pred_loss = loss_fn(targets, predictions)

    gradients = tape.gradient(pred_loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
#--------------------------------------------------------------------------------
    return pred_loss, predictions


with tf.summary.create_file_writer(LOG_DIR).as_default() as writer:
    for epoch in range(5):
        for step, (input_batch, target_batch) in enumerate(dataset):
            total_loss, predictions = train_step(input_batch, target_batch)

            if step == 0:
                tf.summary.trace_export(name="all", step=step, profiler_outdir=LOG_DIR)
            tf.summary.scalar('loss', total_loss, step=step)
            writer.flush()
writer.close()

有一个similar unanswered question OP 无法查看任何图表。

Answer 1

我确信有更好的方法，但我刚刚意识到一个简单的解决方法是只使用现有的 tensorboard 回调逻辑：

tb_callback = tf.keras.callbacks.TensorBoard(LOG_DIR)
tb_callback.set_model(model) # Writes the graph to tensorboard summaries using an internal file writer

如果需要，您可以将自己的摘要写入它使用的同一目录中：tf.summary.create_file_writer(LOG_DIR + '/train')。