Github1.3萬星，迅猛發(fā)展的JAX對比TensorFlow、PyTorch

JAX 是機器學習 (ML) 領(lǐng)域的新生力量，它有望使 ML 編程更加直觀、結(jié)構(gòu)化和簡潔。

在機器學習領(lǐng)域，大家可能對 TensorFlow 和 PyTorch 已經(jīng)耳熟能詳，但除了這兩個框架，一些新生力量也不容小覷，它就是谷歌推出的 JAX。很對研究者對其寄予厚望，希望它可以取代 TensorFlow 等眾多機器學習框架。

JAX 最初由谷歌大腦團隊的 Matt Johnson、Roy Frostig、Dougal Maclaurin 和 Chris Leary 等人發(fā)起。

目前，JAX 在 GitHub 上已累積 13.7K 星。

項目地址：https://github.com/google/jax

迅速發(fā)展的 JAX

JAX 的前身是 Autograd，其借助 Autograd 的更新版本，并且結(jié)合了 XLA，可對 Python 程序與 NumPy 運算執(zhí)行自動微分，支持循環(huán)、分支、遞歸、閉包函數(shù)求導，也可以求三階導數(shù)；依賴于 XLA，JAX 可以在 GPU 和 TPU 上編譯和運行 NumPy 程序；通過 grad，可以支持自動模式反向傳播和正向傳播，且二者可以任意組合成任何順序。

開發(fā) JAX 的出發(fā)點是什么？說到這，就不得不提 NumPy。NumPy 是 Python 中的一個基礎(chǔ)數(shù)值運算庫，被廣泛使用。但是 numpy 不支持 GPU 或其他硬件加速器，也沒有對反向傳播的內(nèi)置支持，此外，Python 本身的速度限制阻礙了 NumPy 使用，所以少有研究者在生產(chǎn)環(huán)境下直接用 numpy 訓練或部署深度學習模型。

在此情況下，出現(xiàn)了眾多的深度學習框架，如 PyTorch、TensorFlow 等。但是 numpy 具有靈活、調(diào)試方便、API 穩(wěn)定等獨特的優(yōu)勢。而 JAX 的主要出發(fā)點就是將 numpy 的以上優(yōu)勢與硬件加速結(jié)合。

目前，基于 JAX 已有很多優(yōu)秀的開源項目，如谷歌的神經(jīng)網(wǎng)絡庫團隊開發(fā)了 Haiku，這是一個面向 Jax 的深度學習代碼庫，通過 Haiku，用戶可以在 Jax 上進行面向?qū)ο箝_發(fā)；又比如 RLax，這是一個基于 Jax 的強化學習庫，用戶使用 RLax 就能進行 Q-learning 模型的搭建和訓練；此外還包括基于 JAX 的深度學習庫 JAXnet，該庫一行代碼就能定義計算圖、可進行 GPU 加速?？梢哉f，在過去幾年中，JAX 掀起了深度學習研究的風暴，推動了科學研究迅速發(fā)展。

JAX 的安裝

如何使用 JAX 呢？首先你需要在 Python 環(huán)境或 Google colab 中安裝 JAX，使用 pip 進行安裝：

$ pip install --upgrade jax jaxlib

注意，上述安裝方式只是支持在 CPU 上運行，如果你想在 GPU 執(zhí)行程序，首先你需要有 CUDA、cuDNN ，然后運行以下命令（確保將 jaxlib 版本映射到 CUDA 版本）：

$ pip install --upgrade jax jaxlib==0.1.61+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

現(xiàn)在將 JAX 與 Numpy 一起導入：

import jax
import jax.numpy as jnp
import numpy as np

JAX 的一些特性

使用 grad() 函數(shù)自動微分：這對深度學習應用非常有用，這樣就可以很容易地運行反向傳播，下面為一個簡單的二次函數(shù)并在點 1.0 上求導的示例：

from jax import grad
def f(x):
  return 3*x**2 + 2*x + 5
def f_prime(x):
  return 6*x +2
grad(f)(1.0)
# DeviceArray(8., dtype=float32)
f_prime(1.0)
# 8.0

jit（Just in time） ：為了利用 XLA 的強大功能，必須將代碼編譯到 XLA 內(nèi)核中。這就是 jit 發(fā)揮作用的地方。要使用 XLA 和 jit，用戶可以使用 jit() 函數(shù)或 @jit 注釋。

from jax import jit
x = np.random.rand(1000,1000)
y = jnp.array(x)
def f(x):
  for _ in range(10):
      x = 0.5*x + 0.1* jnp.sin(x)
  return x
g = jit(f)
%timeit -n 5 -r 5 f(y).block_until_ready()
# 5 loops, best of 5: 10.8 ms per loop
%timeit -n 5 -r 5 g(y).block_until_ready()
# 5 loops, best of 5: 341 μs per loop

pmap：自動將計算分配到所有當前設(shè)備，并處理它們之間的所有通信。JAX 通過 pmap 轉(zhuǎn)換支持大規(guī)模的數(shù)據(jù)并行，從而將單個處理器無法處理的大數(shù)據(jù)進行處理。要檢查可用設(shè)備，可以運行 jax.devices（）：

from jax import pmap
def f(x):
  return jnp.sin(x) + x**2
f(np.arange(4))
#DeviceArray([0.       , 1.841471 , 4.9092975, 9.14112  ], dtype=float32)
pmap(f)(np.arange(4))
#ShardedDeviceArray([0.       , 1.841471 , 4.9092975, 9.14112  ], dtype=float32)

vmap：是一種函數(shù)轉(zhuǎn)換，JAX 通過 vmap 變換提供了自動矢量化算法，大大簡化了這種類型的計算，這使得研究人員在處理新算法時無需再去處理批量化的問題。示例如下：

from jax import vmap
def f(x):
  return jnp.square(x)
f(jnp.arange(10))
#DeviceArray([ 0,  1,  4,  9, 16, 25, 36, 49, 64, 81], dtype=int32)
vmap(f)(jnp.arange(10))
#DeviceArray([ 0,  1,  4,  9, 16, 25, 36, 49, 64, 81], dtype=int32)

TensorFlow vs PyTorch vs Jax

在深度學習領(lǐng)域有幾家巨頭公司，他們所提出的框架被廣大研究者使用。比如谷歌的 TensorFlow、Facebook 的 PyTorch、微軟的 CNTK、亞馬遜 AWS 的 MXnet 等。

每種框架都有其優(yōu)缺點，選擇的時候需要根據(jù)自身需求進行選擇。

我們以 Python 中的 3 個主要深度學習框架——TensorFlow、PyTorch 和 Jax 為例進行比較。這些框架雖然不同，但有兩個共同點：

它們是開源的。這意味著如果庫中存在錯誤，使用者可以在 GitHub 中發(fā)布問題（并修復），此外你也可以在庫中添加自己的功能；

由于全局解釋器鎖，Python 在內(nèi)部運行緩慢。所以這些框架使用 C/C++ 作為后端來處理所有的計算和并行過程。

那么它們的不同體現(xiàn)在哪些方面呢？如下表所示，為 TensorFlow、PyTorch、JAX 三個框架的比較。

TensorFlow

TensorFlow 由谷歌開發(fā)，最初版本可追溯到 2015 年開源的 TensorFlow0.1，之后發(fā)展穩(wěn)定，擁有強大的用戶群體，成為最受歡迎的深度學習框架。但是用戶在使用時，也暴露了 TensorFlow 缺點，例如 API 穩(wěn)定性不足、靜態(tài)計算圖編程復雜等缺陷。因此在 TensorFlow2.0 版本，谷歌將 Keras 納入進來，成為 tf.keras。

目前 TensorFlow 主要特點包括以下：

這是一個非常友好的框架，高級 API-Keras 的可用性使得模型層定義、損失函數(shù)和模型創(chuàng)建變得非常容易；

TensorFlow2.0 帶有 Eager Execution（動態(tài)圖機制），這使得該庫更加用戶友好，并且是對以前版本的重大升級；

Keras 這種高級接口有一定的缺點，由于 TensorFlow 抽象了許多底層機制（只是為了方便最終用戶），這讓研究人員在處理模型方面的自由度更小；

Tensorflow 提供了 TensorBoard，它實際上是 Tensorflow 可視化工具包。它允許研究者可視化損失函數(shù)、模型圖、模型分析等。

PyTorch

PyTorch(Python-Torch) 是來自 Facebook 的機器學習庫。用 TensorFlow 還是 PyTorch？在一年前，這個問題毫無爭議，研究者大部分會選擇 TensorFlow。但現(xiàn)在的情況大不一樣了，使用 PyTorch 的研究者越來越多。PyTorch 的一些最重要的特性包括：

與 TensorFlow 不同，PyTorch 使用動態(tài)類型圖，這意味著執(zhí)行圖是在運行中創(chuàng)建的。它允許我們隨時修改和檢查圖的內(nèi)部結(jié)構(gòu)；

除了用戶友好的高級 API 之外，PyTorch 還包括精心構(gòu)建的低級 API，允許對機器學習模型進行越來越多的控制。我們可以在訓練期間對模型的前向和后向傳遞進行檢查和修改輸出。這被證明對于梯度裁剪和神經(jīng)風格遷移非常有效；

PyTorch 允許用戶擴展代碼，可以輕松添加新的損失函數(shù)和用戶定義的層。PyTorch 的 Autograd 模塊實現(xiàn)了深度學習算法中的反向傳播求導數(shù)，在 Tensor 類上的所有操作， Autograd 都能自動提供微分，簡化了手動計算導數(shù)的復雜過程；

PyTorch 對數(shù)據(jù)并行和 GPU 的使用具有廣泛的支持；

PyTorch 比 TensorFlow 更 Python 化。PyTorch 非常適合 Python 生態(tài)系統(tǒng)，它允許使用 Python 類調(diào)試器工具來調(diào)試 PyTorch 代碼。

JAX 

JAX 是來自 Google 的一個相對較新的機器學習庫。它更像是一個 autograd 庫，可以區(qū)分原生的 python 和 NumPy 代碼。JAX 的一些特性主要包括：

正如官方網(wǎng)站所描述的那樣，JAX 能夠執(zhí)行 Python+NumPy 程序的可組合轉(zhuǎn)換：向量化、JIT 到 GPU/TPU 等等；

與 PyTorch 相比，JAX 最重要的方面是如何計算梯度。在 Torch 中，圖是在前向傳遞期間創(chuàng)建的，梯度在后向傳遞期間計算， 另一方面，在 JAX 中，計算表示為函數(shù)。在函數(shù)上使用 grad() 返回一個梯度函數(shù)，該函數(shù)直接計算給定輸入的函數(shù)梯度；

JAX 是一個 autograd 工具，不建議單獨使用。有各種基于 JAX 的機器學習庫，其中值得注意的是 ObJax、Flax 和 Elegy。由于它們都使用相同的核心并且接口只是 JAX 庫的 wrapper，因此可以將它們放在同一個 bracket 下；

Flax 最初是在 PyTorch 生態(tài)系統(tǒng)下開發(fā)的，更注重使用的靈活性。另一方面，Elegy 受 Keras 啟發(fā)。ObJAX 主要是為以研究為導向的目的而設(shè)計的，它更注重簡單性和可理解性。 

參考鏈接：

https://www.askpython.com/python-modules/tensorflow-vs-pytorch-vs-jax

https://jax.readthedocs.io/en/latest/notebooks/quickstart.html

https://jax.readthedocs.io/en/latest/notebooks/quickstart.html

https://www.zhihu.com/question/306496943/answer/557876584