步履不停：TensorFlow 2.4 新功能一覽！
文 /?Goldie Gadde 和 Nikita Namjoshi，TensorFlow
TensorFlow 2.4 正式發(fā)布！隨著對(duì)分布式訓(xùn)練和混合精度提供更多支持，加入新的 Numpy 前端及用于監(jiān)控和診斷性能瓶頸的工具，這個(gè)版本的亮點(diǎn)在于推出新功能，以及對(duì)性能和擴(kuò)展方面的增強(qiáng)。
tf.distribute 的新增功能
參數(shù)服務(wù)器策略
在版本 2.4 中，實(shí)驗(yàn)性引入了?tf.distribute?模塊的支持，可通過?ParameterServerStrategy?和自定義訓(xùn)練循環(huán)對(duì) Keras 模型進(jìn)行異步訓(xùn)練。與?MultiWorkerMirroredStrategy?一樣，ParameterServerStrategy?是一種多工作器數(shù)據(jù)并行策略；但其梯度更新方式為異步執(zhí)行。?
ParameterServerStrategy
https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/ParameterServerStrategy
參數(shù)服務(wù)器訓(xùn)練集群包含工作節(jié)點(diǎn)和參數(shù)服務(wù)器。系統(tǒng)會(huì)在參數(shù)服務(wù)器上創(chuàng)建變量，然后工作節(jié)點(diǎn)會(huì)在每個(gè)步驟中進(jìn)行讀取和更新。變量的讀取和更新會(huì)在各工作節(jié)點(diǎn)上獨(dú)立進(jìn)行，同時(shí)無需采取任何同步操作。由于工作節(jié)點(diǎn)互不依賴，因此該策略具有工作器容錯(cuò)的優(yōu)勢(shì)，并會(huì)在使用搶占式服務(wù)器時(shí)有所助益。?
如要開始使用此策略，請(qǐng)查閱參數(shù)服務(wù)器訓(xùn)練教程。此教程介紹了如何設(shè)置 ParameterServerStrategy，并說明了如何使用?ClusterCoordinator?類來創(chuàng)建資源、調(diào)度函數(shù)和處理任務(wù)失敗。?
參數(shù)服務(wù)器訓(xùn)練教程
https://tensorflow.google.cn/tutorials/distribute/parameter_server_trainingClusterCoordinator
https://tensorflow.google.cn/api_docs/python/tf/distribute/experimental/coordinator/ClusterCoordinator
多工作節(jié)點(diǎn)鏡像策略
MultiWorkerMirroredStrategy?多工作節(jié)點(diǎn)鏡像策略 ? 已順利度過實(shí)驗(yàn)階段，現(xiàn)已成為穩(wěn)定 API 的組成部分。與單個(gè)工作節(jié)點(diǎn)副本?MirroredStrategy?一樣，MultiWorkerMirroredStrategy?通過同步數(shù)據(jù)并行化實(shí)現(xiàn)分布式訓(xùn)練。但利用?MultiWorkerMirroredStrategy，您可以在多臺(tái)機(jī)器上進(jìn)行訓(xùn)練，且每臺(tái)機(jī)器可以都搭載多個(gè) GPU。?
MultiWorkerMirroredStrategy
https://tensorflow.google.cn/api_docs/python/tf/distribute/MultiWorkerMirroredStrategyMirroredStrategy
https://tensorflow.google.cn/api_docs/python/tf/distribute/MirroredStrategy
?
在同步訓(xùn)練中，每個(gè)工作節(jié)點(diǎn)會(huì)在輸入數(shù)據(jù)的不同片段上計(jì)算正向和反向傳遞次數(shù)，并且在每個(gè)步驟結(jié)束時(shí)匯總梯度。對(duì)于這種稱為 All Reduce 的匯總，?MultiWorkerMirroredStrategy?會(huì)使用集合運(yùn)算保持變量同步。集合運(yùn)算是 TensorFlow 圖表中的單個(gè)算子，可以根據(jù)硬件、網(wǎng)絡(luò)拓?fù)浜蛷埩看笮≡?TensorFlow 運(yùn)行時(shí)中自動(dòng)選擇 All Reduce 算法。集合運(yùn)算還可實(shí)現(xiàn)其他集合運(yùn)算，例如廣播和 All Gather。
集合運(yùn)算
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/collective_ops.py
如要開始使用??MultiWorkerMirroredStrategy，請(qǐng)查閱使用 Keras 進(jìn)行多工作器訓(xùn)練教程，該教程已更新了有關(guān)數(shù)據(jù)集分片、保存/加載使用分布策略訓(xùn)練的模型，以及使用?BackupAndRestore?回調(diào)進(jìn)行故障恢復(fù)的詳細(xì)信息。?
使用 Keras 進(jìn)行多工作器訓(xùn)練
https://tensorflow.google.cn/tutorials/distribute/multi_worker_with_kerasBackupAndRestore
https://tensorflow.google.cn/api_docs/python/tf/keras/callbacks/experimental/BackupAndRestore
如果您不熟悉分布式訓(xùn)練，并希望了解入門方法，或者有興趣在 Google 云端平臺(tái) (GCP) 上進(jìn)行分布式訓(xùn)練，請(qǐng)參閱本博文，以獲取關(guān)于關(guān)鍵概念和步驟的介紹。
Keras 的相關(guān)更新
混合精度
在 TensorFlow 2.4 中，Keras 混合精度 API?已順利度過實(shí)驗(yàn)階段，現(xiàn)已成為穩(wěn)定的 API。大多數(shù) TensorFlow 模型使用的是 float32 dtype；但也存在使用更少內(nèi)存的低精度類型（如 float16）?；旌暇戎冈谕荒Ｐ椭型ㄟ^使用 16 位和 32 位浮點(diǎn)類型，以加快訓(xùn)練速度。該 API 可使模型在 GPU 上性能提高 3 倍，在 TPU 上提高 60%。
Keras 混合精度 API
https://tensorflow.google.cn/api_docs/python/tf/keras/mixed_precision
如要使用混合精度 API，您必須使用 Keras 層和優(yōu)化工具，但無需使用其他 Keras 類，例如模型或損失。如果您對(duì)如何利用此 API 實(shí)現(xiàn)性能優(yōu)化頗有興趣，請(qǐng)查閱混合精度教程。
混合精度教程
https://tensorflow.google.cn/guide/mixed_precision
優(yōu)化工具
此版本支持重構(gòu)?tf.keras.optimizers.Optimizer?類，使?model.fit?或自定義訓(xùn)練循環(huán)的用戶能夠編寫任何適用于優(yōu)化工具的訓(xùn)練代碼?，F(xiàn)所有內(nèi)置的?tf.keras.optimizer.Optimizer?子類均可支持使用?gradient_transformers?和?gradient_aggregator?參數(shù)，您可借此輕松定義自定義梯度轉(zhuǎn)換。?
tf.keras.optimizers.Optimizer
https://tensorflow.google.cn/api_docs/python/tf/keras/optimizers/Optimizer
通過重構(gòu)，您現(xiàn)在可以在編寫自定義訓(xùn)練循環(huán)時(shí)將損失張量直接傳遞給?Optimizer.minimize：
tape = tf.GradientTape()
with tape:
y_pred = model(x, training=True)
loss = loss_fn(y_pred, y_true)
# 如下所示，在使用損失“張量”時(shí)，您可以在“tf.GradientTape”中進(jìn)行傳遞。
optimizer.minimize(loss, model.trainable_variables, tape=tape)
此類更改旨在使?Model.fit?和自定義訓(xùn)練循環(huán)都能擺脫優(yōu)化工具細(xì)節(jié)的限制，從而使您無需修改，即可編寫任何適用于優(yōu)化工具的訓(xùn)練代碼。
函數(shù)式 API 模型構(gòu)建的內(nèi)部改進(jìn)
最后，在 Keras 中，TensorFlow 2.4 可支持對(duì) Keras Functional API 內(nèi)部主要結(jié)構(gòu)的重構(gòu)，從而可降低函數(shù)式模型構(gòu)建的內(nèi)存消耗并簡(jiǎn)化觸發(fā)邏輯。開展此類重構(gòu)操作還能夠確保?TensorFlowOpLayers?行為可預(yù)測(cè)，并可與?CompositeTensor?類型的簽名一起使用。
隆重推出 tf.experimental.numpy
TensorFlow 2.4 以?tf.experimental.numpy?形式，實(shí)驗(yàn)性引入了對(duì)?NumPy API?子集的支持。您可借此模塊，運(yùn)行由 TensorFlow 加速的 NumPy 代碼。由于此 API 基于 TensorFlow 構(gòu)建而成，因此可支持訪問所有 TensorFlow API，與 TensorFlow 實(shí)現(xiàn)無縫互操作，并會(huì)通過編譯和自動(dòng)矢量化開展優(yōu)化。例如，TensorFlow ND 數(shù)組可以與 NumPy 函數(shù)進(jìn)行交互，同樣地，TensorFlow NumPy 函數(shù)也可以接受包括?tf.Tensor?和?np.ndarray?在內(nèi)的不同類型的輸入。
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 代碼計(jì)算梯度
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
https://tensorflow.google.cn/api_docs/python/tf/experimental/numpyNumPy API 實(shí)驗(yàn)性支持
https://github.com/tensorflow/community/blob/master/governance/api-reviews.md#experimental-apis
您可以查閱?TensorFlow 指南上的 NumPy API，了解更多關(guān)于使用此 API 的信息。
TensorFlow 指南上的 NumPy API
https://tensorflow.google.cn/guide/tf_numpy
全新性能分析器工具
TensorFlow Profiler 中的多工作器支持
TensorFlow Profiler?是一套用于評(píng)估 TensorFlow 模型訓(xùn)練性能和資源消耗情況的工具。TensorFlow Profiler 可幫助您了解模型中算子的硬件資源消耗、診斷瓶頸并最終加快訓(xùn)練速度。
TensorFlow Profiler?
https://tensorflow.google.cn/guide/profiler
?
之前版本的TensorFlow Profiler 支持監(jiān)控多 GPU、單主機(jī)訓(xùn)練作業(yè)。在現(xiàn)在 2.4 版本中，您可以分析?MultiWorkerMirroredStrategy?訓(xùn)練作業(yè)的性能。例如，您可以使用采樣模型?API 來執(zhí)行按需分析，并連接到?MultiWorkerMirroredStrategy?工作節(jié)點(diǎn)上正在使用的同一服務(wù)器端口：?
# 在模型運(yùn)行之前啟動(dòng)性能分析器服務(wù)器。
tf.profiler.experimental.server.start(6009)
# 在此處插入模型代碼……
# 例如，您的工作器 IP 地址是 10.0.0.2、10.0.0.3、10.0.0.4，然后您
# 希望執(zhí)行 2 秒鐘的性能分析。性能分析數(shù)據(jù)將
# 保存至 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)
采樣模型
https://tensorflow.google.cn/guide/profiler#sampling_mode
或者，您可以通過向 Capture Profile（捕獲分析結(jié)果）工具提供工作節(jié)點(diǎn)地址來使用 TensorBoard 配置文件插件。
分析完成后，您可以使用新的?Pod Viewer?工具選擇一個(gè)訓(xùn)練步驟，并查閱所有工作節(jié)點(diǎn)的分步時(shí)間類別細(xì)分。
Pod Viewer 工具
https://tensorflow.google.cn/guide/profiler#pod_viewer
有關(guān)如何使用 TensorFlow Profiler 的更多信息，請(qǐng)查閱新發(fā)布的?GPU 性能指南。此指南介紹了您在對(duì)模型訓(xùn)練作業(yè)進(jìn)行性能分析時(shí)可能遇到的常見情況，并提供了調(diào)試工作流程來幫助您優(yōu)化性能，無論您是使用單個(gè) GPU、多個(gè) GPU 還是使用多臺(tái)機(jī)器進(jìn)行訓(xùn)練，均可從中受益。
GPU 性能指南
https://tensorflow.google.cn/guide/gpu_performance_analysis
TFLite Profiler
在 2.4 版本中，您亦可在 Android 中啟用對(duì) TFLite 內(nèi)部結(jié)構(gòu)的跟蹤。現(xiàn)在，您可以使用 Android 版 TFLite Profiler 來識(shí)別性能瓶頸。TFLite 性能評(píng)估指南介紹了如何使用 Android Studio CPU 性能分析器和系統(tǒng)跟蹤應(yīng)用添加跟蹤事件，啟用 TFLite 跟蹤以及捕獲跟蹤。
使用 Android 系統(tǒng)跟蹤應(yīng)用進(jìn)行跟蹤的示例
TFLite 性能評(píng)估指南
https://tensorflow.google.cn/lite/performance/measurement#trace_tensorflow_lite_internals_in_android
提供 GPU 支持的新功能
TensorFlow 2.4 可與 CUDA 11 和 cuDNN 8 一起運(yùn)行，以支持最新上市的 NVIDIA Ampere GPU 架構(gòu)。如需了解 CUDA 11 功能的更多信息，請(qǐng)查閱此?NVIDIA 開發(fā)者博客。
NVIDIA 開發(fā)者博客
https://developer.nvidia.com/blog/cuda-11-features-revealed/
此外，我們亦會(huì)默認(rèn)在搭載 Ampere 的 GPU 上啟用對(duì) TensorFloat-32 的支持。TensorFloat-32（簡(jiǎn)稱為“TF32”）是 NVIDIA Ampere GPU 的一種數(shù)學(xué)模式，可加快令某些 float32 算子（例如矩陣乘法和卷積）在 Ampere GPU 上的運(yùn)行速度，但精度降低。如需了解更多信息，請(qǐng)查閱?tf.config.experimental.enable_tensor_float_32_execution?文檔。?
tf.config.experimental.enable_tensor_float_32_execution
https://tensorflow.google.cn/api_docs/python/tf/config/experimental/enable_tensor_float_32_execution
后續(xù)步驟
請(qǐng)參閱版本說明了解更多信息。如需獲取最新消息，敬請(qǐng)閱讀?TensorFlow 公眾號(hào)(TensorFlow_official)，或訂閱 B 站?Google 中國(guó)?(space.bilibili.com/64169458)。如欲分享您的構(gòu)建成果，請(qǐng)通過?Community Spotlight 計(jì)劃?(goo.gle/TFCS)?向我們提交您的作品。如需提供反饋，請(qǐng)?jiān)?GitHub上提交問題。謝謝！
版本說明
https://github.com/tensorflow/tensorflow/releasesGitHub
https://github.com/tensorflow/tensorflow/issues
推薦閱讀
CondInst：性能和速度均超越Mask RCNN的實(shí)例分割模型
mmdetection最小復(fù)刻版(十一)：概率Anchor分配機(jī)制PAA深入分析
MMDetection新版本V2.7發(fā)布，支持DETR，還有YOLOV4在路上!
無需tricks，知識(shí)蒸餾提升ResNet50在ImageNet上準(zhǔn)確度至80%+
不妨試試MoCo，來替換ImageNet上pretrain模型！
mmdetection最小復(fù)刻版(七)：anchor-base和anchor-free差異分析
mmdetection最小復(fù)刻版(四)：獨(dú)家yolo轉(zhuǎn)化內(nèi)幕
機(jī)器學(xué)習(xí)算法工程師
? ??? ? ? ? ? ? ? ? ? ? ? ??????? ??一個(gè)用心的公眾號(hào)
?
