具有自定义训练循环的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和trace函数

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')。