Jax:有望取代Tensorflow，谷歌出品的又一超高性能机器学习框架
Jax:有望取代Tensorflow，谷歌出品的又一超高性能机器学习框架
在机器学习框架方面，JAX是一个新生事物——尽管Tensorflow的竞争对手从技术上讲已经在2018年后已经很完备，但直到最近JAX才开始在更广泛的机器学习研究社区中获得吸引力。
JAX到底是什么?根据JAX官方介绍:
JAX是NumPy在CPU、GPU和TPU上的版本，具有高性能机器学习研究的强大自动微分(automatic differentiation)能力。
接下来，我们会具体认识JAX。
基础介绍
就像上面说的，JAX是加速器支持的numpy以及大部分scipy功能，带有一些通用机器学习操作的便利函数。
我们举个例子
import jax
import jax.numpy as np
def gpu_backed_hidden_layer(x):
return jax.nn.relu(np.dot(W, x) + b)您可以得到numpy精心设计的API，它从2006年就开始使用了，具有Tensorflow和PyTorch等现代ML工具的性能特征。
JAX还包括通过jax.scipy来支持相当大一部分scipy项目：
from jax.scipy.linalg import svd
singular_vectors, singular_values = svd(x)尽管有加速器支持的numpy + scipy版本已经非常有用，但JAX还有一些其他的妙招。首先让我们看看JAX对自动微分的广泛支持。
自动微分·Autograd
Autograd是一个用于在numpy和原生python代码上高效计算梯度的库。Autograd恰好也是JAX的前身。尽管最初的autograd存储库不再被积极开发，但是在autograd上工作的大部分核心团队已经开始全职从事JAX项目。
就像autograd, JAX允许对一个python函数的输出求导，只需调用grad:
from jax import grad
def hidden_layer(x):
return jax.nn.relu(np.dot(W, x) + b)
grad_hidden_layer = grad(hidden_layer)您还可以通过本机的python控制结构进行区分——而不需要使用tf.cond：
def absolute_value(x)
if x >= 0:
return x
else:
return -x
grad_absolute_value = grad(absolute_value)JAX还支持获取高阶导数——grad函数可以任意连接：
from jax.nn import tanh
# grads all the way down
print(grad(grad(grad(tanh)))(1.0))默认情况下，grad 为您提供了逆向模式梯度——这是计算梯度最常用的模式，它依赖于缓存激活来提高向后传递的效率。反模式差分是计算参数更新最有效的方法。但是，特别是在实现依赖于高阶派生的优化方法时，它并不总是最佳选择。JAX通过jacfwd和jacrev为逆向模式自动差分和正向模式自动差分提供了一流的支持:
from jax import jacfwd, jacrev
hessian_fn = jacfwd(jacrev(fn))除了grad、jacfwd和jacrev之外，JAX还提供了一些实用程序，用于计算函数的线性逼近、定义自定义梯度操作，以及作为其自动微分支持的一部分。
加速线性代数·XLA
XLA (Accelerated Linear Algebra)是一个特定域的线性代数代码编译器，它是JAX将python和numpy表达式转换成加速器支持的操作的基础。
除了允许JAX将python + numpy代码转换为可以在加速器上运行的操作之外(就像我们在第一个示例中看到的那样)，XLA支持还允许JAX将多个操作融合到一个内核中。它在计算图中寻找节点簇，这些节点簇可以被重写以减少计算或中间变量的存储。Tensorflow关于XLA的文档使用以下示例来解释问题可以从XLA编译中受益的实例类型。
def unoptimized_fn(x, y, z):
return np.sum(x + y * z)在没有XLA的情况下运行，这将作为3个独立的内核运行——一个乘法、一个加法和一个加法减法。使用XLA运行时，这变成了一个负责所有这三个方面的内核，不需要存储中间变量，从而节省了时间和内存。
向量化和并行性
虽然Autograd和XLA构成了JAX库的核心，但是还有两个JAX函数脱颖而出。你可以使用jax.vmap和jax.pmap用于向量化和基于spmd(单程序多数据)并行的pmap。
为了说明vmap的优点，我们将返回到我们的简单稠密层的示例，它操作一个由向量x表示的示例。
# convention to distinguish between
# jax.numpy and numpy
import numpy as onp
def hidden_layer(x):
return jax.nn.relu(np.dot(W, x + b)
print(hidden_layer(np.random.randn(128)).shape)
# (128,)我们已经编写了隐含层来获取单个向量输入，但实际上我们几乎总是批量处理输入以利用向量化计算。使用JAX，您可以使用任何接受单个输入的函数，并允许它使用JAX .vmap接受一批输入:
batch_hidden_layer = vmap(hidden_layer)
print(batch_hidden_layer(onp.random.randn(32, 128)).shape)
# (32, 128)它的美妙之处在于，它意味着你或多或少地忽略了模型函数中的批处理维数，并且在你构造模型的时候，在你的头脑中少了一个张量维数。
如果您有几个输入都应该向量化，或者您想沿着轴向量化而不是沿着轴0，您可以使用in_axes参数来指定。
batch_hidden_layer = vmap(hidden_layer, in_axes=(0,))JAX用于SPMD paralellism的实用程序，遵循非常类似的API。如果你有一台4-gpu机器和4个例子，你可以使用pmap在每个设备上运行一个例子。
# first dimension must align with number of XLA-enabled devices
spmd_hidden_layer = pmap(hidden_layer)和往常一样，你可以随心所欲地编写函数:
# hypothetical setup for high-throughput inference
outputs = pmap(vmap(hidden_layer))(onp.random.randn(4, 32, 128))
print(outputs.shape)
# (4, 32, 128)为什么是JAX?
JAX不是因为它都比现有的机器学习框架更加干净,或者因为它是比Tensorflow PyTorch更好地设计的东西,而是因为它能让我们更容易尝试更多的想法以及探索更广泛的空间。