Keras 3.0一統江湖!大更新整合PyTorch、JAX,全球250萬開發者在用了
剛剛,Keras 3.0正式發布!
經過5個月的公開Beta測試,深度學習框架Keras 3.0終于面向所有開發者推出。
全新的Keras 3對Keras代碼庫進行了完全重寫,可以在JAX、TensorFlow和PyTorch上運行,能夠解鎖全新大模型訓練和部署的新功能。
「Keras之父」Fran?ois Chollet在最新版本發布之前,也是做了多次預告。目前,有250+萬的開發者都在使用Keras框架。
重磅消息:我們剛剛發布了 Keras 3.0!
在 JAX、TensorFlow 和 PyTorch 上運行 Keras
使用 XLA 編譯更快地訓練
通過新的 Keras 分發 API 解鎖任意數量的設備和主機的訓練運行
它現在在 PyPI 上上線
開發者甚至可以將Keras用作低級跨框架語言,以開發自定義組件,例如層、模型或指標。
只需一個代碼庫,這些組件便可用在JAX、TensorFlow、PyTorch中的原生工作流。
再次讓Keras成為多后端
最初的Keras可以在Theano、TensorFlow、CNTK,甚至MXNet上運行。
2018年,由于Theano和CNTK已停止開發,TensorFlow似乎成為了唯一可行的選擇,于是,Keras將開發重點放在了TensorFlow上。
而到了今年,情況發生了變化。
根據2023年StackOverflow開發者調查,和2022年Kaggle機器學習和數據科學調查等顯示,
TensorFlow擁有55%到60%的市場份額,是ML在生產領域的首選。
而PyTorch擁有40%到45%的市場份額,是ML在研究領域的首選。
與此同時,JAX雖然市場份額要小得多,但已被Google DeepMind、Midjourney、Cohere等生成式AI領域的頂級參與者所接受。
于是,開發團隊對Keras代碼庫進行了完全重寫,新誕生的Keras 3.0基于模塊化后端架構進行了重構,有能力在任意框架上運行。
同時新的Keras也保證了兼容性,比如在使用TensorFlow后端時,你可以簡單地使用 import keras_core as keras 來替換from tensorflow import keras
——現有的代碼將毫無問題地運行,而且由于 XLA 編譯,通常性能略有提高。
Keras vs. TensorFlow
小編在這里給大家舉一個例子,說明如何從TensorFlow的代碼轉換成Keras的形式。
TensorFlow Core Implementation
Keras implementation
相比之下,我們可以清楚地看到Keras帶來的簡潔性。
TensorFlow可以對每個變量進行更精細的控制,而Keras提供了易用性和快速原型設計的能力。
對于一些開發者來說,Keras省去了開發中的一些麻煩,降低了編程復雜性,節省了時間成本。
Keras 3.0新特性
Keras最大的優勢在于,通過出色的UX、API設計和可調試性可實現高速開發。
而且,它還是一個經過實戰考驗的框架,并為世界上一些最復雜、最大規模的ML系統提供支持,比如Waymo自動駕駛車、YouTube推薦引擎。
那么,使用新的多后端Keras 3還有哪些額外的優勢呢?
- 始終為模型獲得最佳性能。
在基準測試中,發現JAX通常在GPU、TPU和CPU上提供最佳的訓練和推理性能,但結果因模型而異,因為非XLA TensorFlow在GPU上偶爾會更快。
它能夠動態選擇為模型提供最佳性能的后端,而無需對代碼進行任何更改,這意味著開發者可以以最高效率進行訓練和服務。
- 為模型解鎖生態系統可選性。
任何Keras 3模型都可以作為PyTorch模塊實例化,可以作為 TensorFlow SavedModel 導出,也可以作為無狀態 JAX 函數實例化。
這意味著開發者可以將Keras 3模型與PyTorch生態系統包,全系列TensorFlow部署和生產工具(如TF-Serving,TF.js和TFLite)以及JAX大規模TPU訓練基礎架構一起使用。使用 Keras 3 API 編寫一個 model.py ,即可訪問 ML 世界提供的一切。
- 利用JAX的大規模模型并行性和數據并行性。
Keras 3包含一個全新的分布式 API,即keras.distribution 命名空間,目前已在JAX后端實現(即將在TensorFlow和PyTorch后端實現)。
通過它,可以在任意模型尺度和聚類尺度上輕松實現模型并行、數據并行以及兩者的組合。由于它能將模型定義、訓練邏輯和分片配置相互分離,因此使分發工作流易于開發和維護。
- 最大限度地擴大開源模型版本的覆蓋面。
想要發布預訓練模型?想讓盡可能多的人能夠使用它嗎?如果你在純TensorFlow或PyTorch中實現它,它將被大約一半的社區使用。
如果你在Keras 3中實現了它,那么任何人都可以立即使用它,無論他們選擇的框架是什么(即使他們自己不是Keras用戶)。在不增加開發成本的情況下實現2倍的影響。
- 使用來自任何來源的數據管道。
Keras 3 / fit() / evaluate() predict() 例程與 tf.data.Dataset 對象、PyTorch DataLoader 對象、NumPy 數組、Pandas 數據幀兼容——無論你使用什么后端。你可以在 PyTorch DataLoader 上訓練 Keras 3 + TensorFlow 模型,也可以在tf.data.Dataset上訓練Keras 3 + PyTorch模型。
預訓練模型
現在,開發者即可開始使用Keras 3的各種預訓練模型。
所有40個Keras應用程序模型( keras.applications 命名空間)在所有后端都可用。KerasCV和KerasNLP中的大量預訓練模型也適用于所有后端。
其中包括:
- BERT
- OPT
- Whisper
- T5
- Stable Diffusion
- YOLOv8
跨框架開發
Keras 3能夠讓開發者創建在任何框架中都相同的組件(如任意自定義層或預訓練模型),它允許訪問適用于所有后端的 keras.ops 命名空間。
Keras 3包含NumPy API的完整實現,——不是「類似 NumPy」,而是真正意義上的 NumPy API,具有相同的函數和參數。比如 ops.matmul、ops.sum、ops.stack、ops.einsum 等函數。
Keras 3還包含NumPy中沒有的,一組特定于神經網絡的函數,例如 ops.softmax, ops.binary_crossentropy, ops.conv等。
另外,只要開發者使用的運算,全部來自于keras.ops ,那么自定義的層、損失函數、優化器就可以跨越JAX、PyTorch和TensorFlow,使用相同的代碼。
開發者只需要維護一個組件實現,就可以在所有框架中使用它。
Keras架構
下面,我們來稍稍理解一下Keras的機制和架構。
在Keras中,Sequential 和 Model 類是模型構建的核心,為組裝層和定義計算圖提供了一個框架。
Sequential 是層的線性堆棧。它是Model 的子類,專為簡單情況而設計,模型由具有一個輸入和一個輸出的線性層堆棧組成。
Sequential 類有以下一些主要特點:
簡單性:只需按照要執行的順序列出圖層即可。
自動前向傳遞:當向Sequential模型添加層時,Keras會自動將每一層的輸出連接到下一層的輸入,從而創建前向傳遞,而無需手動干預。
內部狀態管理:Sequential管理層的狀態(如權重和偏置)和計算圖。調用compile時,它會通過指定優化器、損失函數和指標來配置學習過程。
訓練和推理:Sequential類提供了fit、evaluate和predict等方法,分別用于訓練、評估和預測模型。這些方法在內部處理訓練循環和推理過程。
Model類與函數式API一起使用,提供了比Sequential更大的靈活性。它專為更復雜的架構而設計,包括具有多個輸入或輸出、共享層和非線性拓撲的模型。
Model 類的主要特點有:
層圖:Model允許創建層圖,允許一個層連接到多個層,而不僅僅是上一個層和下一個層。
顯式輸入和輸出管理:在函數式API中,可以顯式定義模型的輸入和輸出。相比于Sequential,可以允許更復雜的架構。
連接靈活性:Model類可以處理具有分支、多個輸入和輸出以及共享層的模型,使其適用于簡單前饋網絡以外的廣泛應用。
狀態和訓練管理:Model類管理所有層的狀態和訓練過程,同時提供了對層的連接方式,以及數據在模型中的流動方式的更多控制。
Model 類和 Sequential類都依賴于以下機制:
層注冊:在這些模型中添加層時,層會在內部注冊,其參數也會添加到模型的參數列表中。
自動微分:在訓練過程中,Keras使用后端引擎(TensorFlow等)提供的自動微分來計算梯度。這一過程對用戶而言是透明的。
后端執行:實際計算(如矩陣乘法、激活等)由后端引擎處理,后端引擎執行模型定義的計算圖。
序列化和反序列化:這些類包括保存和加載模型的方法,其中涉及模型結構和權重的序列化。
從本質上講,Keras中的Model和Sequential類抽象掉了定義和管理計算圖所涉及的大部分復雜性,使用戶能夠專注于神經網絡的架構,而不是底層的計算機制。
Keras 自動處理各層如何相互連接、數據如何在網絡中流動以及如何進行訓練和推理操作等錯綜復雜的細節。
對于Keras的大更新,有網友使用下面的圖片表達自己的看法:
雖然小編也不知道為什么要炸TensorFlow。
還有網友表示剛好可以用上:
另一位網友發來賀電,「在PyTorch之上使用Keras是一項了不起的成就!」
當然也有網友唱反調,「我想知道為什么有人會使用Keras + Torch而不是普通的 Torch,因為Torch與Tensorflow不同,它有一組很好的API」。
此時Tensorflow的內心:啊對對對,你們說得都對。