Tensorboard Graph with custom training loop does not include my Model

问题

I have created my own loop as shown in the TF 2 migration guide here.
I am currently able to see the graph for only the --- VISIBLE --- section of the code below. How do I make my model (defined in the ---NOT VISIBLE--- section) visible in tensorboard?

If I was not using a custom training loop, I could have gone with the documented model.fit approach:

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

In TF 1, the approach used to be quite straightforward:

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

The Tensorboard migration guide clearly states (here) that:

No direct writing of tf.compat.v1.Graph - instead use @tf.function and trace functions

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()

There's a similar unanswered question where the OP was unable to view any graph.

回答1:

I'm sure there's a better way, but I just realized that a simple workaround is to just use the existing tensorboard callback logic:

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

If you want, you could write your own summaries into the same directory it uses: tf.summary.create_file_writer(LOG_DIR + '/train').

来源：https://stackoverflow.com/questions/61172053/tensorboard-graph-with-custom-training-loop-does-not-include-my-model

标签

python

tensorflow

tensorflow2.0

tensorboard