『JAX中文文档』JAX快速入门

秒速五厘米 2023-01-16 12:24 92阅读 0赞

[![商标][c6bf886a980cc04c7387763e64010eee.png]][c6bf886a980cc04c7387763e64010eee.png 1]

最新的

> 原文： [https://jax.readthedocs.io/en/latest/notebooks/quickstart.html][https_jax.readthedocs.io_en_latest_notebooks_quickstart.html]

# **JAX快速入门** #

--------------------

首先解答一个问题：**JAX是什么？**

简单的说就是GPU加速、支持自动微分(autodiff)的numpy。众所周知，numpy是Python下的基础数值运算库，得到广泛应用。用Python搞科学计算或机器学习，没人离得开它。但是numpy不支持GPU或其他硬件加速器，也没有对backpropagation的内置支持，再加上Python本身的速度限制，所以很少有人会在生产环境下直接用numpy训练或部署深度学习模型。这也是为什么会出现Theano, TensorFlow, Caffe等深度学习框架的原因。但是numpy有其独特的优势：底层、灵活、调试方便、API稳定且为大家所熟悉（与MATLAB一脉相承），深受研究者的青睐。JAX的主要出发点就是将numpy的以上优势与硬件加速结合。现在已经开源的JAX ( [https://github.com/google/jax][https_github.com_google_jax]) 就是通过GPU (CUDA)来实现硬件加速。出自：[https://www.zhihu.com/question/306496943/answer/557876584][https_www.zhihu.com_question_306496943_answer_557876584]

> **小宋说：JAX 其实就是一个支持加速器（GPU 和 TPU）的科学计算库(numpy, scipy)和神经网络库(提供relu，sigmoid, conv 等），相较于PyTorch与TensorFlow更加灵活，通用性更佳。这也是笔者推荐学习和做这个翻译工作的原因，带着大家一起去学习掌握这个框架。**
> 
> 由于笔者非英语专业，有些内荣难免翻译有误，欢迎大家批评指正。对于有些笔者不确定的翻译，采用下划线加括号引用原词的方式来补充，例如：**自动微分****（****differentiation****）**

***官方定义：*JAX是CPU，GPU和TPU上的NumPy，具有出色的自动差分（****differentiation****），可用于高性能机器学习研究。**

作为更新版本的[Autograd][]，JAX可以自动微分本机Python和NumPy代码。它可以通过Python的大部分功能（包括循环，if，递归和闭包）进行微分，甚至可以采用派生类的派生类。它支持反向模式和正向模式微分，并且两者可以任意顺序组成。

新功能是JAX使用 [XLA][] 在诸如GPU和TPU的加速器上编译和运行您的NumPy代码。默认情况下，编译是在后台进行的，而库调用将得到及时的编译和执行。但是，JAX甚至允许您使用单功能API即时将自己的Python函数编译为XLA优化的内核。编译和自动微分可以任意组合，因此您无需离开Python即可表达复杂的算法并获得最佳性能。

import jax.numpy as jnp
    from jax import grad, jit, vmap
    from jax import random

## **乘法矩阵** ##

在以下示例中，我们将生成随机数据。NumPy和JAX之间的一大区别是生成随机数的方式。有关更多详细信息，请参见[JAX中的Common Gotchas][JAX_Common Gotchas]。

key = random.PRNGKey(0)
    x = random.normal(key, (10,))
    print(x)

[-0.372111    0.2642311  -0.18252774 -0.7368198  -0.44030386 -0.15214427
     -0.6713536  -0.59086424  0.73168874  0.56730247]

乘以两个大矩阵。

size = 3000
    x = random.normal(key, (size, size), dtype=jnp.float32)
    %timeit jnp.dot(x, x.T).block_until_ready()  # runs on the GPU

489 ms ± 3.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

我们补充说，`block_until_ready`因为默认情况下JAX使用异步执行（请参见[异步调度][Link 1]）。

JAX NumPy函数可在常规NumPy数组上使用。

import numpy as np
    x = np.random.normal(size=(size, size)).astype(np.float32)
    %timeit jnp.dot(x, x.T).block_until_ready()

488 ms ± 942 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

这样比较慢，因为它每次都必须将数据传输到GPU。您可以使用来确保NDArray由设备内存支持[`device_put()`][device_put]。

from jax import device_put
    
    x = np.random.normal(size=(size, size)).astype(np.float32)
    x = device_put(x)
    %timeit jnp.dot(x, x.T).block_until_ready()

487 ms ± 9.94 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

的输出[`device_put()`][device_put]仍然像NDArray一样，但是它仅在需要打印，绘图，保存（printing, plotting, saving）到磁盘，分支等需要它们的值时才将值复制回CPU。的行为[`device_put()`][device_put]等效于函数，但是速度更快。`jit(lambda x: x)`

如果您有GPU（或TPU！），这些调用将在加速器上运行，并且可能比在CPU上快得多。

x = np.random.normal(size=(size, size)).astype(np.float32)
    %timeit np.dot(x, x.T)

235 ms ± 546 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

JAX不仅仅是一个由GPU支持的NumPy。它还带有一些程序转换，这些转换在编写数字代码时很有用。目前，主要有三个：

*  [`jit()`][jit]，以加快您的代码
 *  [`grad()`][grad]，用于求梯度（derivatives）
 *  [`vmap()`][vmap]，用于自动矢量化或批处理。

让我们一一介绍。我们还将最终以有趣的方式编写这些内容。

## **利用[`jit()`][jit]加快功能** ##

JAX在GPU上透明运行（如果没有，则在CPU上运行，而TPU即将推出！）。但是，在上面的示例中，JAX一次将内核分配给GPU一次操作。如果我们有一系列操作，则可以使用`@jit`装饰器使用[XLA][]一起编译多个操作。让我们尝试一下。

def selu(x, alpha=1.67, lmbda=1.05):
      return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
    
    x = random.normal(key, (1000000,))
    %timeit selu(x).block_until_ready()

4.4 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

我们可以使用加快速度`@jit`，它将在第一次`selu`调用jit-compile并将其之后缓存。

selu_jit = jit(selu)
    %timeit selu_jit(x).block_until_ready()

860 µs ± 27.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

## **通过 `grad()计算梯度`** ##

除了评估数值函数外，我们还希望对其进行转换。一种转变是[自动微分][Link 2]。在JAX中，就像在[Autograd中一样][Autograd 1]，您可以使用[`grad()`][grad]函数来计算梯度。

def sum_logistic(x):
      return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
    
    x_small = jnp.arange(3.)
    derivative_fn = grad(sum_logistic)
    print(derivative_fn(x_small))

[0.25       0.19661197 0.10499357]

让我们以极限微分（finite differences）验证我们的结果是正确的。

def first_finite_differences(f, x):
      eps = 1e-3
      return jnp.array([(f(x + eps * v) - f(x - eps * v)) / (2 * eps)
                       for v in jnp.eye(len(x))])
    
    
    print(first_finite_differences(sum_logistic, x_small))

[0.24998187 0.1964569  0.10502338]

求解梯度可以通过简单调用[`grad()`][grad]。[`grad()`][grad]并[`jit()`][jit]可以任意混合。在上面的示例中，我们先抖动`sum_logistic`然后取其派生词。我们继续深入学习实验：

print(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))

-0.035325594

对于更高级的autodiff，可以将其[`jax.vjp()`][jax.vjp]用于反向模式矢量雅各比积和[`jax.jvp()`][jax.jvp]正向模式雅可比矢量积。两者可以彼此任意组合，也可以与其他JAX转换任意组合。这是组合它们以构成有效计算完整的Hessian矩阵的函数的一种方法：

from jax import jacfwd, jacrev
    def hessian(fun):
      return jit(jacfwd(jacrev(fun)))

## **自动向量化 [`vmap()`][vmap]** ##

JAX在其API中还有另一种转换，您可能会发现它有用：[`vmap()`][vmap]向量化映射。它具有沿数组轴映射函数的熟悉语义（ familiar semantics），但不是将循环保留在外部，而是将循环推入函数的原始操作中以提高性能。当与组合时[`jit()`][jit]，它的速度可以与手动添加批处理尺寸一样快。

我们将使用一个简单的示例，并使用将矩阵向量乘积提升为矩阵矩阵乘积[`vmap()`][vmap]。尽管在这种特定情况下很容易手动完成此操作，但是相同的技术可以应用于更复杂的功能。

mat = random.normal(key, (150, 100))
    batched_x = random.normal(key, (10, 100))
    
    def apply_matrix(v):
      return jnp.dot(mat, v)

给定诸如之类的功能`apply_matrix`，我们可以在Python中循环执行批处理维度，但是这样做的性能通常很差。

def naively_batched_apply_matrix(v_batched):
      return jnp.stack([apply_matrix(v) for v in v_batched])
    
    print('Naively batched')
    %timeit naively_batched_apply_matrix(batched_x).block_until_ready()

Naively batched

4.43 ms ± 9.91 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

我们知道如何手动批处理此操作。在这种情况下，`jnp.dot`透明地处理额外的批次尺寸。

@jit
    def batched_apply_matrix(v_batched):
      return jnp.dot(v_batched, mat.T)
    
    print('Manually batched')
    %timeit batched_apply_matrix(batched_x).block_until_ready()

Manually batched

51.9 µs ± 1.72 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

但是，假设没有批处理支持，我们的功能更加复杂。我们可以用来[`vmap()`][vmap]自动添加批处理支持。

@jit
    def vmap_batched_apply_matrix(v_batched):
      return vmap(apply_matrix)(v_batched)
    
    print('Auto-vectorized with vmap')
    %timeit vmap_batched_apply_matrix(batched_x).block_until_ready()

Auto-vectorized with vmap

79.7 µs ± 249 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

当然，[`vmap()`][vmap]可以与任意组成[`jit()`][jit]，[`grad()`][grad]和任何其它JAX变换。

这只是JAX可以做的事情。我们很高兴看到您的操作！

[c6bf886a980cc04c7387763e64010eee.png]: https://img-blog.csdnimg.cn/img_convert/c6bf886a980cc04c7387763e64010eee.png
[c6bf886a980cc04c7387763e64010eee.png 1]: /images/20221022/23a818ba267c405b8cf7a6a2cd9c8b75.png
[https_jax.readthedocs.io_en_latest_notebooks_quickstart.html]: https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
[https_github.com_google_jax]: https://link.zhihu.com/?target=https%3A//github.com/google/jax
[https_www.zhihu.com_question_306496943_answer_557876584]: https://www.zhihu.com/question/306496943/answer/557876584
[Autograd]: https://github.com/hips/autograd
[XLA]: https://www.tensorflow.org/xla
[JAX_Common Gotchas]: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers
[Link 1]: https://jax.readthedocs.io/en/latest/async_dispatch.html#async-dispatch
[device_put]: https://jax.readthedocs.io/en/latest/jax.html#jax.device_put
[jit]: https://jax.readthedocs.io/en/latest/jax.html#jax.jit
[grad]: https://jax.readthedocs.io/en/latest/jax.html#jax.grad
[vmap]: https://jax.readthedocs.io/en/latest/jax.html#jax.vmap
[Link 2]: https://en.wikipedia.org/wiki/Automatic_differentiation
[Autograd 1]: https://github.com/HIPS/autograd
[jax.vjp]: https://jax.readthedocs.io/en/latest/jax.html#jax.vjp
[jax.jvp]: https://jax.readthedocs.io/en/latest/jax.html#jax.jvp