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

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

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

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

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

項(xiàng)目地址：https://github.com/google/jax

迅速發(fā)展的 JAX

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

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

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

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

JAX 的安裝

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

$ pip install --upgrade jax jaxlib

注意，上述安裝方式只是支持在 CPU 上運(yùn)行，如果你想在 GPU 執(zhí)行程序，首先你需要有 CUDA、cuDNN ，然后運(yùn)行以下命令（確保將 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 一起導(dǎo)入：

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

JAX 的一些特性

使用 grad() 函數(shù)自動(dòng)微分：這對(duì)深度學(xué)習(xí)應(yīng)用非常有用，這樣就可以很容易地運(yùn)行反向傳播，下面為一個(gè)簡單的二次函數(shù)并在點(diǎn) 1.0 上求導(dǎo)的示例：

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 的強(qiáng)大功能，必須將代碼編譯到 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：自動(dòng)將計(jì)算分配到所有當(dāng)前設(shè)備，并處理它們之間的所有通信。JAX 通過 pmap 轉(zhuǎn)換支持大規(guī)模的數(shù)據(jù)并行，從而將單個(gè)處理器無法處理的大數(shù)據(jù)進(jìn)行處理。要檢查可用設(shè)備，可以運(yùn)行 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 變換提供了自動(dòng)矢量化算法，大大簡化了這種類型的計(jì)算，這使得研究人員在處理新算法時(shí)無需再去處理批量化的問題。示例如下：

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

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

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

我們以 Python 中的 3 個(gè)主要深度學(xué)習(xí)框架——TensorFlow、PyTorch 和 Jax 為例進(jìn)行比較。這些框架雖然不同，但有兩個(gè)共同點(diǎn)：

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

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

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

TensorFlow

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

目前 TensorFlow 主要特點(diǎn)包括以下：

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

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

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

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

PyTorch

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

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

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

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

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

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

JAX 

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

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

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

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

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

參考鏈接：

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