TensorFlow 2.9已经发布了!亮点包括oneDNN的性能改进,以及DTensor的发布,这是一个用于模型分布的新API,可以用来无缝地从数据并行转移到模型并行。
我们还对核心库进行了改进,包括Eigen和tf.function
统一,确定性行为,以及对Windows的WSL2的新支持。最后,我们为tf.function retracing和Keras Optimizers发布了新的实验性API。让我们来看看这些新的和改进的功能。
改进的CPU性能:默认为oneDNN
我们与英特尔合作,将oneDNN性能库与TensorFlow整合在一起,以实现英特尔CPU上的顶级性能。自TensorFlow 2.5以来,TensorFlow对oneDNN有实验性的支持,它可以提供高达4倍的性能改进。在TensorFlow 2.9中,我们在Linux x86包上默认开启了oneDNN优化,并为具有神经网络重点硬件功能的CPU开启了优化,如AVX512_VNNI、AVX512_BF16、AMX等,这些功能在英特尔Cascade Lake和更新的CPU上都有。
在启用oneDNN优化的情况下运行TensorFlow的用户可能会观察到与关闭优化时略有不同的数字结果。这是因为浮点舍入方法和顺序不同,可能会产生轻微的错误。如果这给你带来问题,在运行TensorFlow程序之前,通过设置TF_ENABLE_ONEDNN_OPTS=0
,关闭优化。要启用或重新启用它们,在运行你的TensorFlow程序之前设置TF_ENABLE_ONEDNN_OPTS=1
。要验证优化是否开启,在你的程序日志中寻找以"oneDNN custom operations are on"
开始的信息。我们欢迎在GitHub和TensorFlow论坛上的反馈。
使用DTensor的模型并行化
DTensor是一个用于分布式模型处理的新的TensorFlow API,它允许模型从数据并行到基于单程序多数据(SPMD)的模型并行的无缝移动,包括空间分区。这意味着你有工具可以轻松地训练模型,而模型的权重或输入是如此之大,以至于不适合在单个设备上。(如果你熟悉TF1中的Mesh TensorFlow,DTensor也有类似的功能)。
DTensor的设计是以下列原则为核心的:
- 一个与设备无关的API: 这使得相同的模型代码可以在CPU、GPU或TPU上使用,包括跨设备类型的模型分区。
- 多客户端执行:移除协调器,让每个任务驱动其本地连接的设备,允许扩展一个模型而不影响启动时间。
- 全局视角与每个副本:传 统上,TensorFlow的分布式模型代码是围绕副本编写的,但在DTensor中,模型代码是从全局角度编写的,每个副本的代码由DTensor运行时生成和运行。在其他方面,这意味着没有关于批处理规范化是在全局层面还是在每个副本层面发生的不确定性。
我们已经开发了几个关于DTensor的介绍性教程,从DTensor概念到用Keras训练DTensor ML模型:
tf.function的跟踪类型
我们修改了tf.function的回溯方式,使其更加简单、可预测和可配置:
tf.function
的所有参数都被分配了一个tf.types.experimental.TraceType.
自定义用户类可以使用跟踪协议(tf.types.experimental.SupportsTracingProtocol
)声明一个TraceType
。
TraceType
系统使人容易理解追溯规则。例如,子类型化规则表明什么类型的参数可以用于特定的函数追踪。子类型化还解释了不同的特定形状是如何加入到一个通用的形状中的,该形状是它们的超类型,以减少一个函数的痕迹数量。
要了解更多,请看新的API [tf.types.experimental.TraceType](https://www.tensorflow.org/api_docs/python/tf/types/experimental/TraceType)
, [tf.types.experimental.SupportsTracingProtocol](https://www.tensorflow.org/api_docs/python/tf/types/experimental/SupportsTracingProtocol)
,以及tf.function
的reduce_retracing
参数。
对WSL2的支持
Windows Subsystem for Linux让开发者直接在Windows上运行Linux环境,而没有传统的虚拟机或双启动设置的开销。TensorFlow现在支持WSL2开箱即用,包括GPU加速。关于要求和如何在Windows上安装WSL2的更多细节,请参见文档。
确定性的行为
APItf.config.experimental.enable_op_determinism
使TensorFlow操作具有确定性。
确定性意味着,如果你用相同的输入多次运行一个操作,该操作每次都会返回完全相同的输出。这对于调试模型很有用,如果你用确定性从头开始训练你的模型几次,你的模型权重每次都是一样的。通常情况下,由于在操作中使用线程,可以以非确定性的顺序添加浮点数,所以许多操作是非确定性的。
TensorFlow 2.8引入了一个API来使OPS具有确定性,而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优化器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 averaging")。
对于大多数用户来说,你将需要采取任何行动。但是,如果你的高级工作流程属于以下情况,请做相应的修改。
用例1:你在Keras优化器的基础上实现了一个定制的优化器
对于这些工作,请首先检查是否有可能将你的依赖关系改为 [tf.keras.optimizers.experimental.Optimizer.](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/experimental)
如果出于任何原因,你决定继续使用旧的优化器(我们不鼓励这样做),那么你可以将你的优化器改为tf.keras.optimizers.legacy.Optimizer
,以避免在以后的TensorFlow版本中被自动切换到新的优化器。
用例2:您的工作依赖于第三方基于Keras的优化器(如tensorflow_addons)。
只要库继续支持特定的优化器,你的工作就应该成功运行。然而,如果库的维护者未能采取行动来适应Keras优化器的变化,你的工作就会出错。因此,请继续关注第三方库的公告,如果你的工作因优化器故障而出错,请 向Keras团队提交错误。
用例3:你的工作是基于TF1的
首先,请尝试迁移到TF2。这是值得的,而且可能比你想象的要容易得多!tf.keras.optimizers.legacy.XXX
如果由于任何原因,迁移不会很快发生,那么请将您的tf.keras.optimizers.XXX
,以避免被自动切换到新的优化器。
用例4:你的工作有定制的梯度聚合逻辑
通常这意味着你在优化器之外进行梯度聚合,并且用experimental_aggregate_gradients=False
来调用apply_gradients()
。我们改变了参数名称,所以请将你的优化器改为tf.keras.optimizers.experimental.Optimizer
,并设置skip_gradients_aggregation=True
。如果在做了这个改动后出现错误,请向Keras团队提交一个bug。
用例5:你的工作直接调用了已废弃的优化器公共API
请检查你的方法调用在这里是否有匹配。将你的优化器改为tf.keras.optimizers.experimental.Optimizer
。如果出于任何原因,你想继续使用旧的优化器,将你的优化器改为tf.keras.optimizers.legacy.Optimizer
。