Keras 3 效能基準測試
我們對 Keras 3 的三個後端 ( TensorFlow、JAX、PyTorch ) 以及使用 TensorFlow 的 Keras 2 進行了基準測試。請參閱這裡以取得重現結果的程式碼和設定詳細資訊。
模型
我們為生成式和非生成式 AI 任務選擇了一組流行的電腦視覺和自然語言處理模型。 請參閱下表了解我們的選擇。
表 1：基準測試中使用的模型。
| 非生成式 | 生成式 | |
|---|---|---|
| 電腦視覺 (CV) | SegmentAnything1 | StableDiffusion2 |
| 自然語言處理 (NLP) | BERT3 | Gemma4、Mistral5 |
我們衡量的不是每個框架可以達到的最佳效能，而是常見使用者工作流程的開箱即用效能。 為了實現此目標，我們利用了 KerasCV 和 KerasHub 中現有的模型 Keras 版本實作。
硬體
所有基準測試均在 Google Cloud Compute Engine 的單個 NVIDIA A100 GPU 上進行，該 GPU 具有 40GB GPU 記憶體，機器類型為 a2-highgpu-1g，具有 12 個 vCPU 和 85GB 主機記憶體。
結果
表 2 顯示了以每步驟毫秒為單位的基準測試結果。 每個步驟都涉及對單一批次資料進行訓練或預測。 結果為 100 個步驟的平均值，不包括第一個步驟，其中包含模型建立和編譯的開銷。
為了進行公平比較，如果它是相同的模型和任務（擬合或預測），我們會在各框架中使用相同的批次大小。 但是，對於不同的模型和任務，由於它們的大小和架構不同，我們使用不同的批次大小來避免記憶體不足（太大）或 GPU 利用率不足（太小）。
對於大型語言模型（Gemma 和 Mistral），我們也使用了相同的批次大小，因為它們是具有相似參數數量 (7B) 的相同模型類型。 我們還對批次大小為 1 的文字生成進行了基準測試，因為這是使用者廣泛要求的。 我們在訓練和推論中使用 bfloat16 精度，並在訓練（微調）中使用 LoRA6。
為了衡量開箱即用的效能，我們嘗試使用所有預設設定。 例如，使用高階 API（例如，使用 Keras model.fit()），盡可能減少設定。
請注意，這與衡量特定硬體/框架/模型組合的優化實作有很大不同。 請參閱MLPerf 以取得不同框架的最佳優化結果。
表 2：基準測試結果。 速度以毫秒/步驟測量。 越低越好。
| 批次 大小 |
Keras 2 (TensorFlow) |
Keras 3 (TensorFlow) |
Keras 3 (JAX) |
Keras 3 (PyTorch) (即時) |
Keras 3 (最佳) |
|
|---|---|---|---|---|---|---|
| SegmentAnything (擬合) |
1 | 386.93 | 355.25 | 361.69 | 1,388.87 | 355.25 |
| SegmentAnything (預測) |
4 | 1,859.27 | 438.50 | 376.34 | 1,720.96 | 376.34 |
| Stable Diffusion (擬合) |
8 | 1,023.21 | 392.24 | 391.21 | 823.44 | 391.21 |
| Stable Diffusion (預測) |
13 | 649.71 | 616.04 | 627.27 | 1,337.17 | 616.04 |
| BERT (擬合) |
32 | 486.00 | 214.49 | 222.37 | 808.68 | 214.49 |
| BERT (預測) |
256 | 470.12 | 466.01 | 418.72 | 1,865.98 | 418.72 |
| Gemma (擬合) |
8 | 不適用 | 232.52 | 273.67 | 525.15 | 232.52 |
| Gemma (生成) |
32 | 不適用 | 1,134.91 | 1,128.21 | 7,952.67* | 1,128.21 |
| Gemma (生成) |
1 | 不適用 | 758.57 | 703.46 | 7,649.40* | 703.46 |
| Mistral (擬合) |
8 | 不適用 | 185.92 | 213.22 | 452.12 | 185.92 |
| Mistral (生成) |
32 | 不適用 | 966.06 | 957.25 | 10,932.59* | 957.25 |
| Mistral (生成) |
1 | 不適用 | 743.28 | 679.30 | 11,054.67* | 679.30 |
* 目前，使用 PyTorch 後端的 LLM 推論速度異常慢，因為 KerasHub 使用靜態序列填充，這與 HuggingFace 不同。 這將很快解決。
討論
主要發現 1：沒有「最佳」後端
Keras 的三個後端中的每一個都具有獨特的優勢。 至關重要的是，從效能的角度來看，沒有哪個單一後端始終優於其他後端。 最快的後端通常取決於您的特定模型架構。
這突顯了在追求最佳效能時框架可選性的價值。 Keras 3 使您可以無縫切換後端，確保您可以為您的模型找到理想的匹配項。
主要發現 2：Keras 3 比 Keras 2 快
我們還計算了 Keras 3（使用其效能最佳的後端）相較於表 1 中使用 TensorFlow 的 Keras 2 的輸送量（步驟/毫秒）增加量。 結果顯示在下圖中。
圖 1：Keras 3 相較於 Keras 2 的加速，以輸送量（步驟/毫秒）衡量
Keras 3 在所有基準測試模型上的效能都始終優於 Keras 2，在許多情況下速度都大幅提高。 SegmentAnything 推論的速度提升了驚人的 380%，StableDiffusion 訓練輸送量增加了 150% 以上，BERT 訓練輸送量增加了 100% 以上。
重要的是，即使您只是升級到 Keras 3 並繼續使用 TensorFlow 後端，您仍然會看到效能提升。 這主要是因為 Keras 2 直接使用更多 TensorFlow 融合操作，這在某些用例中對於 XLA 編譯可能不是最佳的。
結論
框架效能很大程度上取決於特定模型。 Keras 3 使您可以為您的任務選擇最快的框架 – 這個選項幾乎總是優於 Keras 2。
參考資料
1 Kirillov, Alexander 等人。「Segment anything」。ICCV (2023)。
2 Rombach, Robin 等人。「使用潛在擴散模型的高解析度影像合成」。CVPR (2022)。
3 Kenton, Jacob 等人。「BERT：用於語言理解的深度雙向轉換器的預訓練」。NAACL (2019)。
4 Banks, Jeanine 等人。「Gemma：介紹新的最先進的開放模型」。The Keyword, Google (2024)。
5 Jiang, Albert Q. 等人。「Mistral 7B」。arXiv 預印本 arXiv:2310.06825 (2023)。
6 Hu, Edward J. 等人。「Lora：大型語言模型的低秩適應」。ICLR (2022)。