譯者 | 布加迪
審校 | 重樓
近年來，Keras和Tensorflow組合遇到了一種與之競爭的框架：JAX，它在深度學習開發(fā)者社區(qū)逐漸變得很重要。那么JAX到底是什么？它有哪些功能？它與Keras API又有什么相似和不同之處？Keras API一直是使用Tensorflow（最龐大的Python深度學習庫）的幾乎通用的方法。本文逐一解答了這些問題。
Keras是什么？
Keras誕生于2015年，這種接口用于簡化使用成熟的庫來構建神經(jīng)網(wǎng)絡架構，比如Tensorflow。盡管Keras最初作為一種獨立的框架而創(chuàng)建，但它最終成為了與Tensorflow結合使用的框架：Tensorflow是用于高效訓練和使用可擴展深度神經(jīng)網(wǎng)絡的主要Python庫。隨后，Keras成為Tensorflow上面的抽象層：換句話說，它使“原始”Tensorflow用起來變得容易多了。
Keras便于實現(xiàn)神經(jīng)網(wǎng)絡架構最常見的構建模塊：神經(jīng)元層、目標及激活函數(shù)以及優(yōu)化器等等。特殊類型的深度神經(jīng)網(wǎng)絡架構使用Keras抽象類和方法可以輕松構建，比如卷積神經(jīng)網(wǎng)絡（CNN）和循環(huán)神經(jīng)網(wǎng)絡（RNN）。
JAX是什么？
JAX是一種比較新的框架，不僅適用于深度學習，還適用于整個機器學習開發(fā)。它于2018年由谷歌發(fā)布，側重于高性能數(shù)值計算。具體來說，JAX使Python和numpy（其最大的數(shù)值計算庫）用起來更簡單更快捷，同時無縫支持GPU和TPU高性能處理。就科學計算和數(shù)值計算而言，這是相對普通numpy的一個重要優(yōu)勢，因為numpy只支持CPU執(zhí)行。
由于兼顧高性能執(zhí)行模式的直觀性和多功能，JAX正迅速名聲大噪，成為機器學習和深度學習開發(fā)的最先進框架，有機會最終取代Tensorflow和PyTorch等其他框架。它的自動微分特性有助于高效地執(zhí)行訓練深度神經(jīng)網(wǎng)絡背后的基于梯度的復雜計算。
簡而言之，JAX將科學計算和高性能計算的功能整合到單單一個框架中。
Keras和JAX的異同
現(xiàn)在我們對Keras和JAX已有了大致的了解，下面列出這兩種框架共有的一些特性和諸多不同的方面。
相似之處：
- 深度學習模型開發(fā)：這兩種框架都被廣泛用于構建和訓練深度學習模型。
- GPU/TPU加速：Keras和JAX都可以利用GPU和TPU等加速硬件高效地訓練模型。
- 自動微分：這兩種框架結合了自動計算梯度的機制，梯度計算是模型在訓練過程中優(yōu)化的關鍵過程。
- 與深度學習庫的互操作性：這兩種框架都與流行的深度學習庫TensorFlow兼容。
差異之處：
- 抽象級別：雖然兩種解決方案都提供了一定程度的抽象，但Keras更適合尋求高級API且易于使用的用戶；而JAX更注重控制的靈活性，停留在較低的抽象級別，專注于數(shù)值計算。
- 后端：Keras完全基于并依賴Tensorflow作為后端。同時，JAX不依賴Tensorflow，而是使用一種名為適時（JIT）編譯的方法。話雖如此，JAX和Tensorflow可以結合使用，它們在某些情況下可以很好地互補，比如將高級數(shù)學轉換整合到高級深度學習架構中。
- 易用性：與抽象級別密切相關，Keras旨在易于快速使用。雖然JAX功能更強大，但需要更深入的技術知識才能順利地使用它。
- 函數(shù)轉換：這是JAX獨有的特性，允許高級轉換功能，比如自動向量化和并行執(zhí)行。
- 自動優(yōu)化：同樣，JAX在這方面很突出，它更加靈活，便于在神經(jīng)網(wǎng)絡范圍之外優(yōu)化各種函數(shù)（這就是為什么它也適用于其他機器學習方法，比如集成學習），Keras專門專注于深度學習模型。
我該選擇哪種框架？
了解了這兩種框架之間的異同之后，根據(jù)手頭的問題或場景決定選擇哪種框架就不是什么麻煩事了。
如果用戶尋求易用性、更平緩的學習曲線和更高的抽象級別，Keras是不二的選擇。這個基于Tensorflow庫的API將使用戶能夠在短時間內構建原型，并利用各種深度學習模型處理預測和推理任務。
另一方面，對于經(jīng)驗豐富的開發(fā)人員來說，JAX是一種更強大、更通用的選擇，可以獲得優(yōu)化計算和高級函數(shù)轉換之類的附加功能，而不是嚴格局限于Tensorflow或深度學習建模，不過它需要用戶更大的控制度和低級工程決策。
原文標題：Keras vs. JAX: A Comparison，作者：Iván Palomares Carrascosa
