如何在 TensorFlow 2 中写入 TensorBoard

Question

我对 TensorFlow 1.x 非常熟悉，我正在考虑为即将到来的项目切换到 TensorFlow 2。我在理解如何使用 custom training loop 将标量写入 TensorBoard 日志并通过急切执行时遇到了一些麻烦。

问题描述

在 tf1 中，您将创建一些摘要操作（每个要存储的内容一个操作），然后将其合并为一个操作，在会话中运行该合并的操作，然后使用FileWriter 对象。假设sess 是我们的tf.Session()，下面是一个如何工作的示例：

# While defining our computation graph, define summary ops:
# ... some ops ...
tf.summary.scalar('scalar_1', scalar_1)
# ... some more ops ...
tf.summary.scalar('scalar_2', scalar_2)
# ... etc.

# Merge all these summaries into a single op:
merged = tf.summary.merge_all()

# Define a FileWriter (i.e. an object that writes summaries to files):
writer = tf.summary.FileWriter(log_dir, sess.graph)

# Inside the training loop run the op and write the results to a file:
for i in range(num_iters):
    summary, ... = sess.run([merged, ...], ...)
    writer.add_summary(summary, i)

问题是 tf2 中不再存在会话，我不希望禁用急切执行来完成这项工作。 official documentation 是为 tf1 编写的，我能找到的所有参考资料都建议使用 Tensorboard keras 回调。但是，据我所知，这仅在您通过 model.fit(...) 而不是通过 custom training loop 训练模型时才有效。

我尝试过的

• tf.summary 的 tf1 版本在会话之外运行。显然这些函数的任何组合都失败了，因为 FileWriters、merge_ops 等在 tf2 中甚至都不存在。
• This medium post 表示在一些 tensorflow API 中进行了“清理”，包括 tf.summary()。他们建议使用from tensorflow.python.ops.summary_ops_v2，这似乎不起作用。这个implies 使用record_summaries_every_n_global_steps；稍后会详细介绍。
• 其他一系列帖子1、2、3，建议使用tf.contrib.summary和tf.contrib.FileWriter。但是，tf.contribhas been removed from the core TensorFlow repository and build process。
• A TensorFlow v2 showcase from the official repo，它再次使用 tf.contrib 汇总以及前面提到的 record_summaries_every_n_global_steps。我也无法让它工作（即使不使用 contrib 库）。

tl;博士

我的问题是：

• 有没有办法在 TensroFlow 2 中正确使用tf.summary？
• 如果没有，当使用自定义训练循环（不是model.fit()）时，是否有另一种方法可以在 TensorFlow 2 中编写 TensorBoard 日志？

Answer 1

是的，在 TensorFlow v2 中有一种更简单、更优雅的方式来使用摘要。

首先，创建一个存储日志的文件写入器（例如，在名为 log_dir 的目录中）：

writer = tf.summary.create_file_writer(log_dir)

您想在日志文件中写入内容（例如标量）的任何地方都可以在作者创建的上下文中使用您的旧 tf.summary.scalar。假设您要为步骤i 存储scalar_1 的值：

with writer.as_default():
    tf.summary.scalar('scalar_1', scalar_1, step=i)

您可以在训练循环内部或外部打开任意数量的这些上下文。

例子：

# create the file writer object
writer = tf.summary.create_file_writer(log_dir)

for i, (x, y) in enumerate(train_set):

    with tf.GradientTape() as tape:
        y_ = model(x)
        loss = loss_func(y, y_)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

    # write the loss value
    with writer.as_default():
        tf.summary.scalar('training loss', loss, step=i+1)