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

·  阅读 678

商标

最新的

原文: jax.readthedocs.io/en/latest/n…

JAX快速入门


首先解答一个问题:JAX是什么?

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

 

小宋说: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

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使用异步执行(请参见异步调度)。

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

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()仍然像NDArray一样,但是它仅在需要打印,绘图,保存(printing, plotting, saving)到磁盘,分支等需要它们的值时才将值复制回CPU。的行为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(),以加快您的代码
  • grad(),用于求梯度(derivatives)
  • vmap(),用于自动矢量化或批处理。

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

 

利用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()计算梯度

除了评估数值函数外,我们还希望对其进行转换。一种转变是自动微分。在JAX中,就像在Autograd中一样,您可以使用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()jit()可以任意混合。在上面的示例中,我们先抖动sum_logistic然后取其派生词。我们继续深入学习实验:

print(grad(jit(grad(jit(grad(sum_logistic)))))(1.0))
复制代码
-0.035325594
复制代码

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

from jax import jacfwd, jacrev
def hessian(fun):
  return jit(jacfwd(jacrev(fun)))
复制代码

自动向量化 vmap()

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

我们将使用一个简单的示例,并使用将矩阵向量乘积提升为矩阵矩阵乘积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()自动添加批处理支持。

@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()可以与任意组成jit()grad()和任何其它JAX变换。

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

收藏成功!
已添加到「」, 点击更改