TensorFlow 2.9 新增功能

tfwechat · June 27, 2022, 6:44am

发布人：TensorFlow 团队的 Goldie Gadde 和 Douglas Yarrington

TensorFlow 2.9 现已发布！其中亮点包括 oneDNN 性能改进，以及新发布的 DTensor（DTensor 是一种新的模型分布 API，可将数据并行无缝迁移至模型并行）。

我们还改进了核心库，包括 Eigen 和 tf.function 统一、确定性行为以及对 Windows WSL2 的新支持。此外，我们还针对 tf.function retracing 和 Keras 优化器发布了新的实验性 API。

下面让我们来详细了解 TensorFlow 2.9 的全新改进内容。

提高 CPU 性能：默认启用 oneDNN

我们与 Intel 合作，将 oneDNN 性能库与 TensorFlow 进行集成，以便在 Intel CPU 上实现更好的性能。自 TensorFlow 2.5 版本发布以来，TensorFlow 便增加了对 oneDNN 的实验性支持，实现了高达 4 倍的性能提升。

在 TensorFlow 2.9 中，我们为 Linux x86 软件包、具有神经网络硬件功能（如 AVX512_VNNI、AVX512_BF16、AMX 等）的 CPU，以及 Intel Cascade Lake 和其他新版 CPU 默认启用 oneDNN 优化。

浮点舍入方法和顺序不同可能会产生轻微的误差，因此运行 TensorFlow 时启用 oneDNN 优化会出现与关闭优化时略有不同的数值结果。如果上述情况给您造成不便，请在运行 TensorFlow 程序前通过设置 TF_ENABLE_ONEDNN_OPTS=0 来关闭优化。如要启用优化或在关闭后重新启用，请在运行 TensorFlow 程序前设置 TF_ENABLE_ONEDNN_OPTS=1。您可在程序日志中查找以 "oneDNN custom operations are on" 开头的消息，以确定优化是否启用。我们期待您在 GitHub 和 TensorFlow 论坛上提出反馈。

使用 DTensor 的模型并行

DTensor 是一种新的 TensorFlow API，可用于分布式模型处理，用户可使用该 API 将模型从数据并行无缝迁移至基于单程序多数据 (SPMD) 的模型并行，包括空间分区。这意味着您可以借助此工具，轻松训练因模型权重或输入因太大而无法放在单个设备上的模型。（如您熟悉的 TF1 中的 Mesh TensorFlow，DTensor 也具有类似功能）

DTensor 的核心设计原则如下：

**与设备无关的 API：**支持在 CPU、GPU 或 TPU 上使用相同的模型代码，包括跨设备类型划分的模型。
**多客户端执行：**移除协调器并让每个任务驱动其本地连接的设备，从而实现在不影响启动时间的情况下扩展模型。
**全局视角 vs. 每个副本：**传统方式中，TensorFlow 的分布式模型代码围绕副本而编写，但在 DTensor 中，模型代码以全局视角编写，而每个副本的代码则由 DTensor 运行时生成和运行。除其他好处外，这样做还意味着用户现可确定批标准化是处于全局级别或每个副本级别。

我们开发了一些 DTensor 入门教程，包含从 DTensor 概念到使用 Keras 训练 DTensor ML 模型的一些列内容：

tf.function 的 TraceType

新版本修改了 tf.function 的回溯方式，使其更简单、可预测和可配置。

tf.function 的所有参数都分配了一个 tf.types.experimental.TraceType. 。自定义用户类可以使用跟踪协议 (tf.types.experimental.SupportsTracingProtocol) 声明 TraceType。

借助 TraceType 系统，用户可以轻松理解追溯规则。例如，子类型 (subtyping) 规则表示可以与特定函数跟踪一起使用的参数类型。子类型规则还可以解释如何将不同的特定形状联接至作为其超类型的通用形状中，以减少函数的跟踪数量。

如要了解更多信息，请参阅 tf.types.experimental.TraceType、tf.types.experimental.SupportsTracingProtocol 和 tf.function 的 reduce_retracing 参数的新 API。

对 WSL2 的支持

开发者可借助 Windows Subsystem for Linux 直接在 Windows 上运行 Linux 环境，而无需传统的虚拟机或双启动设置。TensorFlow 现支持开箱即用的 WSL2，包括 GPU 加速。请参阅文档，详细了解相关要求以及如何在 Windows 上安装 WSL2。

确定性行为

API tf.config.experimental.enable_op_determinism 使 TensorFlow 算子具有确定性。

这表示，如果您使用相同的输入多次运行同一个算子，则该算子每次都会返回完全相同的输出。确定性对调试模型而言很有用，如果借助确定性从头开始多次训练模型，则模型的权重每次都会相同。通常而言，由于在算子中使用了线程（线程会以不确定性的顺序添加浮点数），许多算子都具有不确定性。

我们在 TensorFlow 2.8 版本中推出了一个可使算子具有确定性的 API，而 TensorFlow 2.9 可以在某些情况下提高 tf.data 的确定性性能。如要在运行 TensorFlow 模型时启用确定性，您只需在程序开头添加以下内容：

tf.keras.utils.set_random_seed(1)

tf.config.experimental.enable_op_determinism()

第一行设置 Python、NumPy 和 TensorFlow 的随机种子，这一步是实现确定性的必要操作。第二行使每个 TensorFlow 算子都具有确定性。请注意，启用确定性通常会牺牲性能，因此在启用算子确定性时，模型的运行速度可能会减慢。

利用 Keras 优化训练

我们在 TensorFlow 2.9 版本中发布了一个 Keras Optimizer API 的全新实验性版本，即 tf.keras.optimizers.experimental。此 API 可提供更统一和扩展的内置优化器目录，从而使用户更轻松地自定义和扩展。

在后续版本中，tf.keras.optimizers.experimental.Optimizer（及子类）将取代 tf.keras.optimizers.Optimizer（及子类），这意味着使用旧版 Keras 优化器的工作流将自动切换至使用新版优化器。当前（旧版）tf.keras.optimizers.* API 仍可通过 tf.keras.optimizers.legacy.* 进行访问，例如 tf.keras.optimizers.legacy.Adam。

以下为新优化器类的一些亮点：

部分模型的训练速度逐步加快。

更易于编写自定义优化器。

对模型权重移动平均（“Polyak 平均”）的内置支持。

对大多数用户而言，无需进行任何操作。但是，如果您的高级工作流属于以下情况，请进行相应的更改：

用例 1：基于 Keras 优化器实现自定义优化器

若属于此情况，请您首先检查是否可以将依赖项更改为 tf.keras.optimizers.experimental.Optimizer ，如果因为一些原因您决定继续使用旧版优化器（不建议），那么您可以将优化器更改为 tf.keras.optimizers.legacy.Optimizer，以避免在更高版本的 TensorFlow 中自动切换至新的优化器。

用例 2：工作流依赖于第三方基于 Keras 的优化器（例如 tensorflow_addons）

只要库继续支持特定优化器，工作流就能成功运行。但是，如果库维护者未能采取措施以适应 Keras 优化器的更改，工作流就会出错。因此请持续关注第三方库的公告，如果工作流因优化器故障而中断，请向 Keras 团队提交错误报告。

用例 3：工作流基于 TF1

首先，请尝试迁移至 TF2。不要小瞧这一步，效果可能出人意料！如果因为一些原因您无法立即迁移，那么请将 tf.keras.optimizers.XXX 替换为 tf.keras.optimizers.legacy.XXX 以避免自动切换至新的优化器。

用例 4：工作流具有自定义梯度聚合逻辑

这通常意味着您在优化器之外进行梯度聚合，并使用 experimental_aggregate_gradients=False 调用 apply_gradients()。我们更改了参数名称，因此请将您的优化器更改为 tf.keras.optimizers.experimental.Optimizer 并设置 skip_gradients_aggregation=True。如果更改后仍出现错误，请向 Keras 团队提交错误报告。

用例 5：工作流直接调用已弃用的优化器公开 API

点击此处检查您的方法调用是否属于此情况。将优化器更改为 tf.keras.optimizers.experimental.Optimizer。如果因为一些原因您想继续使用旧版优化器，请将优化器更改为 tf.keras.optimizers.legacy.Optimizer。