TensorFlow 2.4来了：上线对分布式训练和混合精度的新功能支持

欢迎关注“

计算机视觉研究院

”

计算机视觉研究院专栏

计算机视觉研究院

专注CV领域的公众号

作者：Edison

长按二维码关注

今天，谷歌正式发布了TensorFlow 2.4，带来了多项新特性和功能改进。

TensorFlow2.4 的更新包括对于分布式训练和混合精度的新功能支持，对 NumPy API 子集的试验性支持以及一些用于监测性能瓶颈的新工具。

根据TensorFlow官方博客，本次更新的主要内容整理如下：

tf.distribute 中的新功能

参数服务器策略

在TensorFlow2.4中，tf.distribute 模块引入了对使用 ParameterServerStrategy 和自定义训练循环进行模型异步训练的试验性支持。和 MultiWorkerMirroredStrategy 类似，ParameterServerStrategy 是一种多工作器（multi-worker）数据并行策略，但梯度更新是异步的。

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

多工作器镜像策略

MultiWorkerMirroredStrategy 已经脱离试验阶段，成为稳定的 API。像单工作器的 MirroredStrategy 一样，MultiWorkerMirroredStrategy 通过同步数据并行实现分布式训练，顾名思义，借助 MultiWorkerMirroredStrategy 可以在多台机器上进行训练，每台机器都可能具有多个 GPU。

Keras 更新

混合精度

在TensorFlow2.4 中，Keras 混合精度 API 已经脱离试验阶段，成为稳定的 API。大多数TensorFlow模型使用 float32 dtype，但现在有些低精度数据类型占用的内存更少，比如 float16。混合精度指的是在同一模型中使用 16 位和 32 位浮点数以进行更快的训练。这一 API 可将模型性能在 GPU 上提高到 3 倍，在 TPU 上提高 60%。要使用混合精度 API，必须使用 Keras 层和优化器，但不一定需要使用其他 Keras 类。

优化器

本次更新包括重构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 和自定义训练循环与优化器细节更加不相关，从而让使用者无需修改即可编写出与任何优化器共同使用的训练代码。

最后，TensorFlow2.4的更新还包括 Keras Functional API 内部的重构，改善了函数式模型构造所产生的内存消耗并简化了触发逻辑。这种重构可以保证TensorFlowOpLayers的行为可预测，并且可以使用CompositeTensor类型签名（type signature）。

tf.experimental.numpy

TensorFlow2.4引入了对NumPy API子集的试验性支持。该模块可以运行由TensorFlow加速的 NumPy 代码，由于这一 API 是基于TensorFlow构建的，因此可与TensorFlow无缝衔接，允许访问所有TensorFlow API 并通过编译和自动矢量化提供优化后的运行。

例如，TensorFlow ND数组可以与NumPy函数互通，类似地，TensorFlow NumPy函数可以接受包括tf.Tensor和np.ndarray在内的不同类型输入。

import tensorflow.experimental.numpy as tn
# 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 Profiler是度量TensorFlow模型的训练性能和资源消耗情况的工具，用来诊断性能瓶颈，最终加快训练速度。

此前，TensorFlow Profiler支持多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)

另外，你可以通过向捕获配置文件工具提供工作器地址来使用TensorBoard 配置文件插件。配置之后，你可以使用新的 Pod Viewer tool 来选择训练步骤，并查看所有工作器上该步骤的 step-time 的细分。

TFLite Profiler

TFLite Profiler 则支持在 Android 中跟踪 TFLite 内部信息，以识别性能瓶颈。

GPU 支持

TensorFlow 2.4 与 CUDA 11 和 cuDNN 8 配合运行，支持最新发布的英伟达安培架构 GPU，对于 CUDA 11 的新特性，可以参考英伟达开发者博客：

https://developer.nvidia.com/blog/cuda-11-features-revealed/

在新版本中，默认情况下会启用安培 GPU 的新特性——对 TensorFloat-32 的支持。TensorFloat-32 又简称为 TF32，是英伟达 Ampere GPU 的一种数学精度模式，可导致某些 float32 运算（如矩阵乘法和卷积）在安培架构 GPU 上运行得更快，但精度略有降低。

／End.

我们开创“计算机视觉协会”知识星球一年有余，也得到很多同学的认可，我们定时会推送实践型内容与大家分享，在星球里的同学可以随时提问，随时提需求，我们都会及时给予回复及给出对应的答复。

如果想加入我们“计算机视觉研究院”，请扫二维码加入我们。我们会按照你的需求将你拉入对应的学习群！

计算机视觉研究院主要涉及深度学习领域，主要致力于人脸检测、人脸识别，多目标检测、目标跟踪、图像分割等研究方向。研究院接下来会不断分享最新的论文算法新框架，我们这次改革不同点就是，我们要着重”研究“。之后我们会针对相应领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！

微信公众号

公众号:计算机视觉研究院

知乎专栏 : EdisonGzq

******************

点我打印二维码