TensorFlow 2.9上線：oneDNN改進實現CPU性能優化，WSL2開箱即用
昨日，TensorFlow 官方宣布：TensorFlow 2.9 來了！距離上次 2.8 版本 的更新僅僅過去三個月。
新版本亮點包括如下：
- oneDNN 的性能改進；
- DTensor 的發布，這是一種新 API，可用于從數據并行無縫遷移到模型并行；
- 對核心庫進行了改進，包括 Eigen、tf.function 統一以及對 Windows 的 WSL2 的新支持；
- 還為 tf.function retracing 和 Keras 優化器發布了新的實驗性 API。
下面讓我們來看看 TensorFlow 2.9 具體的改進都有哪些。
TensorFlow 2.9 新特性
提高 CPU 性能：oneDNN
TensorFlow 已經與英特爾合作，將 oneDNN 性能庫與 TensorFlow 集成，以實現在英特爾 CPU 上的最高性能。從 TensorFlow 2.5 以來，TensorFlow 已經對 oneDNN 進行了實驗性支持，它可以提供高達 4 倍的性能提升。在 TensorFlow 2.9 中，默認在 Linux x86 包和具有神經網絡硬件功能（如 AVX512_VNNI、AVX512_BF16、AMX 等）的 CPU 上啟用 oneDNN 優化，這些功能可在 Intel Cascade Lake 和更新的 CPU 上找到。
英特爾表示：我們很高興與 GoogleAI 合作，將 oneDNN 庫作為 TensorFlow 2.9 中的默認后端 CPU 優化！TensorFlow 用戶現在將看到由我們的開放軟件優化驅動的 AI 性能加速，無需更改代碼！
在啟用 oneDNN 優化的情況下運行 TensorFlow 的用戶，可能會觀察到與關閉優化時略有不同的數值結果，這是因為浮點舍入方法和順序不同，可能會產生輕微的誤差。
如果上述情況給你帶來麻煩，請在運行 TensorFlow 程序之前通過設置 TF_ENABLE_ONEDNN_OPTS=0 來關閉優化。假如你要重新啟用它們，請在運行 TensorFlow 程序之前設置 TF_ENABLE_ONEDNN_OPTS=1。要驗證優化是否啟用，請在程序日志中查找以 oneDNN custom operations are on 開頭的消息。
DTensor
DTensor 是一種新的 TensorFlow API，可用于分布式模型處理，它允許模型無縫地從數據并行遷移到基于單程序多數據（SPMD）的模型并行。
DTensor 的核心設計原則如下：
- 設備無關 API：這允許在 CPU、GPU 或 TPU 上使用相同的模型代碼，包括跨設備類型劃分的模型；
- 多客戶端執行：移除 coordinator 并讓每個任務驅動其本地連接的設備，從而允許在不影響啟動時間的情況下擴展模型；
- global perspective VS per-replica：傳統上使用 TensorFlow，分布式模型代碼是圍繞 replicas 編寫的，但使用 DTensor，模型代碼是從 global perspective 編寫的，每個 replica 代碼由 DTensor 運行時生成和運行。
TensorFlow 官方已經開發了幾個關于 DTensor 的入門教程，參考資料如下：
- DTensor 概念：https://www.tensorflow.org/guide/dtensor_overview
- 帶有 DTensors 分布式 ML：https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial
- 將 DTensors、Keras 聯合使用：https://www.tensorflow.org/tutorials/distribute/dtensor_keras_tutorial
tf.function 的 TraceType
新版本已經改進了 tf.function 回溯（retraces）方式，使其更簡單、可預測和可配置。
所有關于 tf.function 的參數被分配一個 tf.types.experimental.TraceType。自定義用戶類可以使用跟蹤協議（tf.types.experimental.SupportsTracingProtocol）聲明一個 TraceType。TraceType 系統使理解追溯規則變得很容易。
支持 WSL2
WSL2 允許開發人員直接在 Windows 上運行 Linux 環境，而不需要傳統虛擬機或雙啟動設置。TensorFlow 現在支持 WSL2 開箱即用，包括 GPU 加速。
確定性行為
API tf.config.experimental.enable_op_determinism 使得 TensorFlow 的 op 是確定性的。確定性意味著如果用戶使用相同的輸入多次運行一個 op，則 op 每次都返回完全相同的輸出，這對于調試模型很有用。
通常來講，許多 op 是不確定的，因為在 op 中使用了線程，這些線程可以以不確定的順序添加浮點數。
TensorFlow 2.8 引入了一個 API 來使 op 具有確定性，TensorFlow 2.9 在某些情況下提高了確定性性能。如果你希望 TensorFlow 模型確定性地運行，只需將以下內容添加到程序的開頭：
<code>tf.keras.utils.set_random_seed(1)</code>
<code>tf.config.experimental.enable_op_determinism()</code>
第一行設置 Python、NumPy 和 TensorFlow 的隨機種子，這是確定性所必需的。第二行使每個 TensorFlow op 具有確定性。請注意，確定性通常是以降低性能為代價的，因此當啟用 op 確定性時，你的模型可能會運行得更慢。
使用 Keras 優化訓練
TensorFlow 2.9 中添加了 Keras Optimizer API 的新實驗版本，即 tf.keras.optimizers.experimental。該 API 提供更統一和擴展的內置優化器目錄，用戶可以更輕松地定制和擴展。
在未來的版本中，tf.keras.optimizers.experimental.Optimizer（及子類）將取代 tf.keras.optimizers.Optimizer（及子類），這意味著使用舊版本 Keras 優化器的工作流將自動切換至新的優化器。
如下為新優化器類的一些亮點：
- 一些模型的訓練速度越來越快；
- 更易于編寫自定義優化器；
- 內置對模型權重移動平均的支持（Polyak 平均）。
