TensorFlow 团队的 Goldie Gadde 和 Nikita Namjoshi 发布

TF 2.4 已发布！本版本新增了对分布式训练和混合精度的支持，并提供了新的 NumPy 前端和用于监控和诊断瓶颈的工具，旨在提供性能和扩展性方面的全新功能和增强功能。

tf.distribute 中的新功能

在 2.4 中，tf.distribute 模块引入了对使用 ParameterServerStrategy 和自定义训练循环进行模型异步训练的实验性支持。与 MultiWorkerMirroredStrategy 一样，ParameterServerStrategy 也是一种多工作节点数据并行策略；但是，梯度更新是异步的。

参数服务器训练集群由工作节点和参数服务器组成。变量是在参数服务器上创建的，然后在每个步骤中由工作节点读取和更新。变量的读取和更新在所有工作节点上独立进行，无需任何同步。由于工作节点之间不相互依赖，因此该策略具有工作节点容错的优势，如果您使用的是可抢占式虚拟机，则很有用。

要开始使用此策略，请查看 参数服务器训练教程。本教程将向您展示如何设置 ParameterServerStrategy 并定义训练步骤，以及如何使用 ClusterCoordinator 类将训练步骤的执行调度到远程工作节点。

MultiWorkerMirroredStrategy 已退出实验阶段，现已成为稳定 API 的一部分。与单工作节点对应项 MirroredStrategy 一样，MultiWorkerMirroredStrategy 使用同步数据并行实现分布式训练。但是，顾名思义，使用 MultiWorkerMirroredStrategy，您可以在多台机器上进行训练，每台机器可能包含多个 GPU。

在同步训练中，每个工作节点对输入数据的不同切片计算前向和反向传递，并且在更新模型之前聚合梯度。对于这种称为 全约简 的聚合，MultiWorkerMirroredStrategy 使用 CollectiveOps 来保持变量同步。集体操作是 TensorFlow 图中的单个操作，它可以根据硬件、网络拓扑和张量大小自动在 TensorFlow 运行时选择全约简算法。

要开始使用 MultiWorkerMirroredStrategy，请查看 使用 Keras 进行多工作节点训练 教程，该教程已更新了有关数据集分片、保存/加载使用分布式策略训练的模型以及使用 BackupAndRestore 回调进行故障恢复的详细信息。

如果您不熟悉分布式训练并想了解如何开始，或者您有兴趣在 GCP 上进行分布式训练，请参阅此博文，了解关键概念和步骤的介绍。

Keras 中的更新

在 TensorFlow 2.4 中，Keras 混合精度 API 已退出实验阶段，现已成为稳定 API。大多数 TensorFlow 模型使用 float32 数据类型；但是，还有一些精度较低的类型，例如 float16，它们使用的内存更少。混合精度是在同一模型中使用 16 位和 32 位浮点类型来加快训练速度。此 API 可以将 GPU 上的模型性能提高 3 倍，将 TPU 上的模型性能提高 60%。

要使用混合精度 API，您必须使用 Keras 层和优化器，但无需使用其他 Keras 类，例如模型或损失函数。如果您想了解如何利用此 API 来提高性能，请查看 混合精度教程。

本版本包括对 tf.keras.optimizers.Optimizer 类的重构，使用户能够使用 model.fit 或自定义训练循环编写与任何优化器兼容的训练代码。 所有内置的 tf.keras.optimizer.Optimizer 子类现在都接受 gradient_transformers 和 gradient_aggregator 参数，允许您轻松定义自定义梯度转换。

使用重构后的代码，您现在可以将损失张量直接传递给 Optimizer.minimize，以编写自定义训练循环

tape = tf.GradientTape()
with tape:
  y_pred = model(x, training=True)
  loss = loss_fn(y_pred, y_true)

# You can pass in the `tf.GradientTape` when using a loss `Tensor` as shown below.

optimizer.minimize(loss, model.trainable_variables, tape=tape)

这些更改旨在使 Model.fit 和自定义训练循环更加独立于优化器详细信息，从而使您能够编写与任何优化器兼容的训练代码，无需修改。

最后，TensorFlow 2.4 包括对 Keras 函数式 API 内部机制的重大重构，改进了函数式模型构造的内存消耗，并简化了触发逻辑。此重构还确保 TensorFlowOpLayers 的行为可预测，并且可以使用 CompositeTensor 类型签名。

介绍 tf.experimental.numpy

TensorFlow 2.4 引入了对 NumPy API 子集的 实验性 支持，可作为 tf.experimental.numpy 使用。此模块使您能够运行 NumPy 代码，并通过 TensorFlow 加速。由于它是在 TensorFlow 之上构建的，因此此 API 与 TensorFlow 无缝互操作，允许访问所有 TensorFlow API，并使用编译和自动矢量化提供优化的执行。例如，TensorFlow ND 数组可以与 NumPy 函数互操作，类似地，TensorFlow NumPy 函数可以接受不同类型的输入，包括 tf.Tensor 和 np.ndarray。

import tensorflow.experimental.numpy as tnp

# Use NumPy code in input pipelines

dataset = tf.data.Dataset.from_tensor_slices(
    tnp.random.randn(1000, 1024)).map(
    lambda z: z.clip(-1,1)).batch(100)

# Compute gradients through NumPy code

def grad(x, wt):
  with tf.GradientTape() as tape:
    tape.watch(wt)
    output = tnp.dot(x, wt)
    output = tf.sigmoid(output)
  return tape.gradient(tnp.sum(output), wt)

您可以在 TensorFlow 上的 NumPy API 指南 中详细了解如何使用此 API。

新的性能分析器工具

TensorFlow 性能分析器 是一套工具，您可以使用这些工具来衡量 TensorFlow 模型的训练性能和资源消耗。TensorFlow 性能分析器可以帮助您了解模型中操作的硬件资源消耗，诊断瓶颈，并最终加快训练速度。

以前，TensorFlow 性能分析器支持监控多 GPU、单主机训练作业。在 2.4 中，您现在可以分析 MultiWorkerMirroredStrategy 训练作业。例如，您可以使用 采样模式 API 进行按需性能分析，并连接到 MultiWorkerMirroredStrategy 工作节点使用的同一服务器:端口

# Start a profiler server before your model runs.


tf.profiler.experimental.server.start(6009)

# Model code goes here....
 
# E.g. your worker IP addresses are 10.0.0.2, 10.0.0.3, 10.0.0.4, and you
# would like to profile for a duration of 2 seconds. The profiling data will
# be saved to the Google Cloud Storage path “your_tb_logdir”.
 
tf.profiler.experimental.client.trace(
    'grpc://10.0.0.2:6009,grpc://10.0.0.3:6009,grpc://10.0.0.4:6009',
    'gs://your_tb_logdir',
    2000)

或者，您可以通过将工作节点地址提供给 Capture Profile 工具来使用 TensorBoard 性能分析插件。

完成性能分析后，您可以使用新的 Pod 查看器工具 选择一个训练步骤并查看其在所有工作节点上的步骤时间类别细分。

有关如何使用 TensorFlow 性能分析器的更多信息，请查看新发布的 GPU 性能指南。本指南展示了在分析模型训练作业时可能会遇到的常见情况，并提供了一个调试工作流程，以帮助您提高性能，无论您是在使用一个 GPU、多个 GPU 还是多台机器进行训练。

TFLite 性能分析器可以跟踪 Android 中的 TFLite 内部机制，以识别性能瓶颈。TFLite 性能测量指南 将向您展示如何添加跟踪事件、启用 TFLite 跟踪以及使用 Android Studio CPU 性能分析器和系统跟踪应用捕获跟踪。

使用 Android 系统跟踪应用的示例跟踪

GPU 支持的新功能

TensorFlow 2.4 与 CUDA 11 和 cuDNN 8 一起运行，支持新发布的 NVIDIA Ampere GPU 架构。要详细了解 CUDA 11 功能，请查看此 NVIDIA 开发人员博客。

此外，对 Ampere 架构 GPU 上的 TensorFloat-32 的支持默认情况下已启用。TensorFloat-32（简称 `TF32`）是 NVIDIA Ampere GPU 的一种数学模式，它会导致某些 float32 操作（如矩阵乘法和卷积）在 Ampere GPU 上运行速度更快，但精度降低。要详细了解，请参阅 tf.config.experimental.enable_tensor_float_32_execution 的文档。

后续步骤

查看 发行说明，获取更多信息。要保持最新状态，您可以阅读 TensorFlow 博客，关注 twitter.com/tensorflow，或订阅 youtube.com/tensorflow。如果您构建了一些内容，希望与大家分享，请将其提交到我们的社区亮点区 goo.gle/TFCS。如有反馈，请在 GitHub 上提交问题。感谢您！