譯者 | 布加迪

審校 | 重樓

近年來，Keras和Tensorflow組合遇到了一種與之競爭的框架：JAX，它在深度學習開發者社區逐漸變得很重要。那么JAX到底是什么？它有哪些功能？它與Keras API又有什么相似和不同之處？Keras API一直是使用Tensorflow（最龐大的Python深度學習庫）的幾乎通用的方法。本文逐一解答了這些問題。

Keras是什么？

Keras誕生于2015年，這種接口用于簡化使用成熟的庫來構建神經網絡架構，比如Tensorflow。盡管Keras最初作為一種獨立的框架而創建，但它最終成為了與Tensorflow結合使用的框架：Tensorflow是用于高效訓練和使用可擴展深度神經網絡的主要Python庫。隨后，Keras成為Tensorflow上面的抽象層：換句話說，它使“原始”Tensorflow用起來變得容易多了。

Keras便于實現神經網絡架構最常見的構建模塊：神經元層、目標及激活函數以及優化器等等。特殊類型的深度神經網絡架構使用Keras抽象類和方法可以輕松構建，比如卷積神經網絡（CNN）和循環神經網絡（RNN）。

JAX是什么？

JAX是一種比較新的框架，不僅適用于深度學習，還適用于整個機器學習開發。它于2018年由谷歌發布，側重于高性能數值計算。具體來說，JAX使Python和numpy（其最大的數值計算庫）用起來更簡單更快捷，同時無縫支持GPU和TPU高性能處理。就科學計算和數值計算而言，這是相對普通numpy的一個重要優勢，因為numpy只支持CPU執行。

由于兼顧高性能執行模式的直觀性和多功能，JAX正迅速名聲大噪，成為機器學習和深度學習開發的最先進框架，有機會最終取代Tensorflow和PyTorch等其他框架。它的自動微分特性有助于高效地執行訓練深度神經網絡背后的基于梯度的復雜計算。

簡而言之，JAX將科學計算和高性能計算的功能整合到單單一個框架中。

Keras和JAX的異同

現在我們對Keras和JAX已有了大致的了解，下面列出這兩種框架共有的一些特性和諸多不同的方面。

相似之處：

• 深度學習模型開發：這兩種框架都被廣泛用于構建和訓練深度學習模型。
• GPU/TPU加速：Keras和JAX都可以利用GPU和TPU等加速硬件高效地訓練模型。
• 自動微分：這兩種框架結合了自動計算梯度的機制，梯度計算是模型在訓練過程中優化的關鍵過程。
• 與深度學習庫的互操作性：這兩種框架都與流行的深度學習庫TensorFlow兼容。

差異之處：

• 抽象級別：雖然兩種解決方案都提供了一定程度的抽象，但Keras更適合尋求高級API且易于使用的用戶；而JAX更注重控制的靈活性，停留在較低的抽象級別，專注于數值計算。
• 后端：Keras完全基于并依賴Tensorflow作為后端。同時，JAX不依賴Tensorflow，而是使用一種名為適時（JIT）編譯的方法。話雖如此，JAX和Tensorflow可以結合使用，它們在某些情況下可以很好地互補，比如將高級數學轉換整合到高級深度學習架構中。
• 易用性：與抽象級別密切相關，Keras旨在易于快速使用。雖然JAX功能更強大，但需要更深入的技術知識才能順利地使用它。
• 函數轉換：這是JAX獨有的特性，允許高級轉換功能，比如自動向量化和并行執行。
• 自動優化：同樣，JAX在這方面很突出，它更加靈活，便于在神經網絡范圍之外優化各種函數（這就是為什么它也適用于其他機器學習方法，比如集成學習），Keras專門專注于深度學習模型。

我該選擇哪種框架？

了解了這兩種框架之間的異同之后，根據手頭的問題或場景決定選擇哪種框架就不是什么麻煩事了。

如果用戶尋求易用性、更平緩的學習曲線和更高的抽象級別，Keras是不二的選擇。這個基于Tensorflow庫的API將使用戶能夠在短時間內構建原型，并利用各種深度學習模型處理預測和推理任務。

另一方面，對于經驗豐富的開發人員來說，JAX是一種更強大、更通用的選擇，可以獲得優化計算和高級函數轉換之類的附加功能，而不是嚴格局限于Tensorflow或深度學習建模，不過它需要用戶更大的控制度和低級工程決策。

原文標題：Keras vs. JAX: A Comparison，作者：Iván Palomares Carrascosa