互联网资讯 / 人工智能 · 2023年11月26日 0

TensorFlow 2.4新功能介绍

TensoRFlow 2.4 已正式发布!该版本在分布式训练和混合精度方面提供了更全面的支持,同时引入了新的 NumPy 前端和用于监控及诊断性能瓶颈的工具。这次更新的重点在于新增功能以及性能和扩展能力的提升。

tf.distribute 的新特性

参数服务器策略

在 2.4 版本中,tf.distribute 模块的支持被实验性引入,用户可以通过 ParameterServerStrategy 和自定义训练循环对 Keras 模型进行异步训练。与 MultiWorkerMirroredStrategy 一样,ParameterServerStrategy 也是一种多工作器数据并行策略,但其梯度更新过程为异步执行。

参数服务器训练集群由工作节点和参数服务器组成。系统会在参数服务器上创建变量,工作节点在每个步骤中读取和更新这些变量。变量的读取和更新在各工作节点上独立进行,无需任何同步操作。由于工作节点之间互不依赖,该策略在工作器容错方面具有优势,并在使用抢占式服务器时表现良好。

要开始使用此策略,请参考参数服务器训练教程,该教程详细介绍了如何配置 ParameterServerStrategy,并说明了如何使用 ClusterCoordinator 类来创建资源、调度功能和处理任务失败。

多工作节点镜像策略

MultiWorkerMirroredStrategy 的多工作节点镜像策略已成功进入稳定 API 阶段。与单工作节点副本 MirroredStrategy 相似,MultiWorkerMirroredStrategy 通过同步数据并行化实现分布式训练。使用 MultiWorkerMirroredStrategy,您可以在多台机器上进行训练,每台机器可以配备多个 GPU。

在同步训练中,每个工作节点会在不同的数据片段上执行前向和反向传递,并在每个步骤结束时汇总梯度。MultiWorkerMirroredStrategy 使用集合运算保持变量同步,集合运算是 TensorFlow 图表中的单个算子,可以根据硬件配置、网络拓扑和张量大小在 TensorFlow 运行时中自动选择 AllReduce 算法。此外,集合运算还支持其他集合运算,如广播和 AllGather。

如要开始使用 MultiWorkerMirroredStrategy,请参考 Keras 多工作器训练教程,该教程已更新,涵盖数据集分片、保存/加载使用分布式策略训练的模型以及使用 backupAndRestore 回调进行故障恢复的详细信息。

Keras 的更新

混合精度

在 TensorFlow 2.4 中,Keras 混合精度 API 已成功进入稳定阶段。大多数 TensorFlow 模型使用 float32 数据类型,但也支持使用更少内存的低精度类型(如 float16)。混合精度指在同一模型中通过使用 16 位和 32 位浮点类型来加速训练。该 API 能使模型在 GPU 上性能提升 3 倍,在 TPU 上提升 60%。

要使用混合精度 API,您需要采用 Keras 层和优化工具,但不需要使用其他 Keras 类,例如模型或损失。如果您想了解如何利用此 API 实现性能优化,请参考混合精度教程。

优化工具

此次版本支持对 tf.keras.optimizers.Optimizer 类进行重构,允许用户在 Model.fit 或自定义训练循环中编写适用于优化工具的训练代码。现有所有内置的 tf.keras.optimizers.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) # 如下所示,在使用损失“张量”时,您可以在“tf.GradientTape”中进行传递。 optimizer.minimize(loss, Model.trainable_variables, tape=tape)

这些更改旨在使 Model.fit 和自定义训练循环不再受限于优化工具的细节,从而使您无需修改即可编写任何适用于优化工具的训练代码。

函数式 API 模型构建的内部改进

在 Keras 中,TensorFlow 2.4 支持对 Keras Functional API 主要结构的重构,从而降低函数式模型构建的内存消耗并简化触发逻辑。此类重构操作确保了 TensorFlow 操作行为的可预测性,并可与 CompositeTensor 类型的签名配合使用。

全新推出的 tf.experimental.numpy

TensorFlow 2.4 实验性引入了 tf.experimental.numpy,对 NumPy API 子集的支持。借助此模块,您可以运行由 TensorFlow 加速的 NumPy 代码。由于此 API 基于 TensorFlow 构建,因此支持访问所有 TensorFlow API,并实现无缝互操作,同时通过编译和自动矢量化进行优化。例如,TensorFlow ND 数组可以与 NumPy 函数进行交互,反之亦然,TensorFlow NumPy 函数也可以接受不同类型的输入,包括 tf.Tensor 和 np.ndarray。

import tensorflow.experimental.numpy as tnp # 在输入流水线中使用 NumPy 代码 dataset = tf.data.Dataset.from_tensor_slices( tnp.random.randn(1000, 1024)).map( lambda z: z.clip(-1, 1)).batch(100) # 通过 NumPy 代码计算梯度 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) tf.experimental.numpy

有关 NumPy API 实验性支持的更多信息,您可以查阅 TensorFlow 指南。

全新性能分析工具

TensorFlow Profiler 中的多工作器支持

TensorFlow Profiler 是一套用于评估 TensorFlow 模型训练性能和资源消耗的工具。它帮助您了解模型中算子的硬件资源消耗、诊断瓶颈并加速训练。

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

# 在模型运行之前启动性能分析器服务器。 tf.Profiler.experimental.server.start(6009) # 在此处插入模型代码 # 例如,您的工作器 IP 地址是 10.0.0.2、10.0.0.3、10.0.0.4,然后您 # 希望执行 2 秒钟的性能分析。性能分析数据将 # 保存至 Google Cloud Storage 路径“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 Viewer 工具选择一个训练步骤,并查看所有工作节点的分步时间类别细分。

Pod Viewer 工具

有关如何使用 TensorFlow Profiler 的更多信息,请查阅新发布的 GPU 性能指南。该指南介绍了您在对模型训练作业进行性能分析时可能遇到的常见情况,并提供了调试工作流程,帮助您优化性能,无论您是使用单个 GPU、多个 GPU 还是多台机器进行训练,均可从中受益。

TFLite Profiler

在 2.4 版本中,您还可以在 Android 中启用对 TFLite 内部结构的跟踪。现在,您可以使用 Android 版 TFLite Profiler 来识别性能瓶颈。TFLite 性能评估指南介绍了如何使用 Android Studio CPU 性能分析器和系统跟踪应用添加跟踪事件,启用 TFLite 跟踪以及捕获跟踪。

使用 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 算子的运行速度(例如矩阵乘法和卷积),但精度稍有降低。如需了解更多信息,请查阅 tf.config.experimental.enable_tensor_float_32_execution 文档。