将 tf.summary 用法迁移到 TF 2.x

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本
import tensorflow as tf
2022-12-14 12:08:16.683364: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 12:08:16.683474: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 12:08:16.683485: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

TensorFlow 2.x 对用于写入 TensorBoard 可视化摘要数据的 tf.summary API 进行了重大更改。

发生了哪些变化

tf.summary API 视为两个子 API 很有用

在 TF 1.x 中

这两个部分必须手动连接在一起 - 通过 Session.run() 获取摘要操作输出,并调用 FileWriter.add_summary(output, step)v1.summary.merge_all() 操作通过使用图集合来聚合所有摘要操作输出,从而简化了此过程,但这种方法在急切执行和控制流方面仍然效果不佳,使其特别不适合 TF 2.x。

在 TF 2.X 中

这两个部分紧密集成在一起,现在单个 tf.summary 操作在执行时会立即写入其数据。从模型代码中使用 API 应该仍然看起来很熟悉,但现在它对急切执行很友好,同时仍然与图模式兼容。将 API 的两个部分集成在一起意味着 summary.FileWriter 现在是 TensorFlow 执行上下文的一部分,并由 tf.summary 操作直接访问,因此配置写入器是看起来不同的主要部分。

急切执行的示例用法(TF 2.x 中的默认设置)

writer = tf.summary.create_file_writer("/tmp/mylogs/eager")

with writer.as_default():
  for step in range(100):
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)
    writer.flush()
ls /tmp/mylogs/eager
events.out.tfevents.1671019701.kokoro-gcp-ubuntu-prod-1817880111.7700.0.v2

使用 tf.function 图执行的示例用法

writer = tf.summary.create_file_writer("/tmp/mylogs/tf_function")

@tf.function
def my_func(step):
  with writer.as_default():
    # other model code would go here
    tf.summary.scalar("my_metric", 0.5, step=step)

for step in tf.range(100, dtype=tf.int64):
  my_func(step)
  writer.flush()
ls /tmp/mylogs/tf_function
events.out.tfevents.1671019701.kokoro-gcp-ubuntu-prod-1817880111.7700.1.v2

使用传统 TF 1.x 图执行的示例用法

g = tf.compat.v1.Graph()
with g.as_default():
  step = tf.Variable(0, dtype=tf.int64)
  step_update = step.assign_add(1)
  writer = tf.summary.create_file_writer("/tmp/mylogs/session")
  with writer.as_default():
    tf.summary.scalar("my_metric", 0.5, step=step)
  all_summary_ops = tf.compat.v1.summary.all_v2_summary_ops()
  writer_flush = writer.flush()


with tf.compat.v1.Session(graph=g) as sess:
  sess.run([writer.init(), step.initializer])

  for i in range(100):
    sess.run(all_summary_ops)
    sess.run(step_update)
    sess.run(writer_flush)
ls /tmp/mylogs/session
events.out.tfevents.1671019702.kokoro-gcp-ubuntu-prod-1817880111.7700.2.v2

转换您的代码

将现有的 tf.summary 用法转换为 TF 2.x API 无法可靠地自动执行,因此 tf_upgrade_v2 脚本 只会将其全部重写为 tf.compat.v1.summary,并且不会自动启用 TF 2.x 行为。

部分迁移

为了使迁移到 TF 2.x 更加容易,对于仍然严重依赖 TF 1.x 摘要 API 记录操作(如 tf.compat.v1.summary.scalar())的模型代码用户来说,可以先只迁移写入器 API,以便稍后完全迁移模型代码中的单个 TF 1.x 摘要操作。

为了支持这种迁移方式,tf.compat.v1.summary 会在以下条件下自动转发到其 TF 2.x 等效项

请注意,当调用 TF 2.x 摘要实现时,返回值将是一个空字节串张量,以避免重复写入摘要。此外,输入参数转发是尽力而为的,并非所有参数都会保留(例如,family 参数将得到支持,而 collections 参数将被删除)。

tf.compat.v1.summary.scalar 中调用 tf.summary.scalar 行为的示例

# Enable eager execution.
tf.compat.v1.enable_v2_behavior()

# A default TF 2.x summary writer is available.
writer = tf.summary.create_file_writer("/tmp/mylogs/enable_v2_in_v1")
# A step is set for the writer.
with writer.as_default(step=0):
  # Below invokes `tf.summary.scalar`, and the return value is an empty bytestring.
  tf.compat.v1.summary.scalar('float', tf.constant(1.0), family="family")

完全迁移

为了完全迁移到 TF 2.x,您需要按照以下步骤调整代码

  1. 必须存在通过 .as_default() 设置的默认写入器才能使用摘要操作

    • 这意味着以急切方式执行操作或在图构建中使用操作
    • 如果没有默认写入器,摘要操作将变成静默的无操作
    • 默认写入器不会(尚未)跨越 @tf.function 执行边界传播 - 它们仅在函数被追踪时被检测到 - 所以最佳实践是在函数体中调用 writer.as_default(),并确保写入器对象在 @tf.function 被使用时一直存在
  2. 必须通过 step 参数将“步长”值传递到每个操作中

    • TensorBoard 需要步长值才能将数据渲染为时间序列
    • 显式传递是必要的,因为 TF 1.x 中的全局步长已被移除,因此每个操作都必须知道要读取的所需步长变量
    • 为了减少样板代码,实验性地支持注册默认步长值,可以通过 tf.summary.experimental.set_step() 使用,但这是一种临时的功能,可能会在未经通知的情况下更改
  3. 单个摘要操作的函数签名已更改

    • 返回值现在是一个布尔值(指示是否实际写入了摘要)
    • 第二个参数名称(如果使用)已从 tensor 更改为 data
    • collections 参数已被移除;集合仅适用于 TF 1.x
    • family 参数已被移除;只需使用 tf.name_scope()
  4. [仅适用于传统的图模式/会话执行用户]

    • 首先使用 v1.Session.run(writer.init()) 初始化写入器

    • 使用 v1.summary.all_v2_summary_ops() 获取当前图的所有 TF 2.x 摘要操作,例如通过 Session.run() 执行它们

    • 使用 v1.Session.run(writer.flush()) 刷新写入器,类似地使用 close()

如果您的 TF 1.x 代码使用的是 tf.contrib.summary API,它与 TF 2.x API 非常相似,因此 tf_upgrade_v2 脚本将自动执行大多数迁移步骤(并针对无法完全迁移的任何用法发出警告或错误)。在大多数情况下,它只是将 API 调用重写为 tf.compat.v2.summary;如果您只需要与 TF 2.x 兼容,您可以删除 compat.v2,并将其引用为 tf.summary.

其他提示

除了上述关键区域之外,一些辅助方面也发生了变化

  • 条件记录(如“每 100 步记录一次”)有了新的外观

    • 要控制操作和相关代码,请将它们包装在常规的 if 语句中(在急切模式下和在 @tf.function 通过 autograph 中有效)或 tf.cond
    • 要仅控制摘要,请使用新的 tf.summary.record_if() 上下文管理器,并将您选择的布尔条件传递给它
    • 这些替换了 TF 1.x 模式

      if condition:
        writer.add_summary()
      
  • 不直接写入 tf.compat.v1.Graph - 而是使用追踪函数

  • 不再使用 tf.summary.FileWriterCache 对每个日志目录进行全局写入器缓存

    • 用户应该要么实现自己的写入器对象缓存/共享,要么只使用单独的写入器(TensorBoard 对后者的支持 正在进行中
  • 事件文件二进制表示已更改

    • TensorBoard 1.x 已经支持新格式;此差异仅影响手动解析事件文件中的摘要数据的用户
    • 摘要数据现在存储为张量字节;您可以使用 tf.make_ndarray(event.summary.value[0].tensor) 将其转换为 numpy