[[416349]]

在機(jī)器學(xué)習(xí)領(lǐng)域，大家可能對(duì) TensorFlow 和 PyTorch 已經(jīng)耳熟能詳，但除了這兩個(gè)框架，一些新生力量也不容小覷，它就是谷歌推出的 JAX。很多研究者對(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 程序；通過(guò) grad，可以支持自動(dòng)模式反向傳播和正向傳播，且二者可以任意組合成任何順序。

開(kāi)發(fā) JAX 的出發(fā)點(diǎn)是什么？說(shuō)到這，就不得不提 NumPy。NumPy 是 Python 中的一個(gè)基礎(chǔ)數(shù)值運(yùn)算庫(kù)，被廣泛使用。但是 numpy 不支持 GPU 或其他硬件加速器，也沒(méi)有對(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)秀的開(kāi)源項(xiàng)目，如谷歌的神經(jīng)網(wǎng)絡(luò)庫(kù)團(tuán)隊(duì)開(kāi)發(fā)了 Haiku，這是一個(gè)面向 Jax 的深度學(xué)習(xí)代碼庫(kù)，通過(guò) Haiku，用戶可以在 Jax 上進(jìn)行面向?qū)ο箝_(kāi)發(fā)；又比如 RLax，這是一個(gè)基于 Jax 的強(qiáng)化學(xué)習(xí)庫(kù)，用戶使用 RLax 就能進(jìn)行 Q-learning 模型的搭建和訓(xùn)練；此外還包括基于 JAX 的深度學(xué)習(xí)庫(kù) JAXnet，該庫(kù)一行代碼就能定義計(jì)算圖、可進(jìn)行 GPU 加速。可以說(shuō)，在過(guò)去幾年中，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è)簡(jiǎn)單的二次函數(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 通過(guò) pmap 轉(zhuǎn)換支持大規(guī)模的數(shù)據(jù)并行，從而將單個(gè)處理器無(wú)法處理的大數(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 通過(guò) vmap 變換提供了自動(dòng)矢量化算法，大大簡(jiǎn)化了這種類型的計(jì)算，這使得研究人員在處理新算法時(shí)無(wú)需再去處理批量化的問(wèn)題。示例如下：

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)：

• 它們是開(kāi)源的。這意味著如果庫(kù)中存在錯(cuò)誤，使用者可以在 GitHub 中發(fā)布問(wèn)題（并修復(fù)），此外你也可以在庫(kù)中添加自己的功能；
• 由于全局解釋器鎖，Python 在內(nèi)部運(yùn)行緩慢。所以這些框架使用 C/C++ 作為后端來(lái)處理所有的計(jì)算和并行過(guò)程。

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

TensorFlow

TensorFlow 由谷歌開(kāi)發(fā)，最初版本可追溯到 2015 年開(kāi)源的 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)來(lái)，成為 tf.keras。

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

• 這是一個(gè)非常友好的框架，高級(jí) API-Keras 的可用性使得模型層定義、損失函數(shù)和模型創(chuàng)建變得非常容易；
• TensorFlow2.0 帶有 Eager Execution（動(dòng)態(tài)圖機(jī)制），這使得該庫(kù)更加用戶友好，并且是對(duì)以前版本的重大升級(jí)；
• Keras 這種高級(jí)接口有一定的缺點(diǎn)，由于 TensorFlow 抽象了許多底層機(jī)制（只是為了方便最終用戶），這讓研究人員在處理模型方面的自由度更小；
• Tensorflow 提供了 TensorBoard，它實(shí)際上是 Tensorflow 可視化工具包。它允許研究者可視化損失函數(shù)、模型圖、模型分析等。

PyTorch

PyTorch(Python-Torch) 是來(lái)自 Facebook 的機(jī)器學(xué)習(xí)庫(kù)。用 TensorFlow 還是 PyTorch？在一年前，這個(gè)問(wèn)題毫無(wú)爭(zhēng)議，研究者大部分會(huì)選擇 TensorFlow。但現(xiàn)在的情況大不一樣了，使用 PyTorch 的研究者越來(lái)越多。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)行越來(lái)越多的控制。我們可以在訓(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)提供微分，簡(jiǎn)化了手動(dòng)計(jì)算導(dǎo)數(shù)的復(fù)雜過(guò)程；
• PyTorch 對(duì)數(shù)據(jù)并行和 GPU 的使用具有廣泛的支持；
• PyTorch 比 TensorFlow 更 Python 化。PyTorch 非常適合 Python 生態(tài)系統(tǒng)，它允許使用 Python 類調(diào)試器工具來(lái)調(diào)試 PyTorch 代碼。

JAX

JAX 是來(lái)自 Google 的一個(gè)相對(duì)較新的機(jī)器學(xué)習(xí)庫(kù)。它更像是一個(gè) autograd 庫(kù)，可以區(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í)庫(kù)，其中值得注意的是 ObJax、Flax 和 Elegy。由于它們都使用相同的核心并且接口只是 JAX 庫(kù)的 wrapper，因此可以將它們放在同一個(gè) bracket 下；
• Flax 最初是在 PyTorch 生態(tài)系統(tǒng)下開(kāi)發(fā)的，更注重使用的靈活性。另一方面，Elegy 受 Keras 啟發(fā)。ObJAX 主要是為以研究為導(dǎo)向的目的而設(shè)計(jì)的，它更注重簡(jiǎn)單性和可理解性。