vlambda博客
学习文章列表

JAX 是机器学习 (ML) 领域的新生力量，它有望使 ML 编程更加直观、结构化和简洁。

$ pip install --upgrade jax jaxlib

$ pip install --upgrade jax jaxlib==0.1.61+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

import jaximport jax.numpy as jnpimport numpy as np

from jax import graddef f(x): return 3*x**2 + 2*x + 5def f_prime(x): return 6*x +2grad(f)(1.0)# DeviceArray(8., dtype=float32)f_prime(1.0)# 8.0

from jax import jitx = 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 xg = 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

from jax import pmapdef f(x): return jnp.sin(x) + x**2f(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)

from jax import vmapdef 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)

• 它们是开源的。这意味着如果库中存在错误，使用者可以在 GitHub 中发布问题（并修复），此外你也可以在库中添加自己的功能；

• 由于全局解释器锁，Python 在内部运行缓慢。所以这些框架使用 C/C++ 作为后端来处理所有的计算和并行过程。

• 这是一个非常友好的框架，高级 API-Keras 的可用性使得模型层定义、损失函数和模型创建变得非常容易；

• TensorFlow2.0 带有 Eager Execution（动态图机制），这使得该库更加用户友好，并且是对以前版本的重大升级；

• Keras 这种高级接口有一定的缺点，由于 TensorFlow 抽象了许多底层机制（只是为了方便最终用户），这让研究人员在处理模型方面的自由度更小；

• Tensorflow 提供了 TensorBoard，它实际上是 Tensorflow 可视化工具包。它允许研究者可视化损失函数、模型图、模型分析等。

• 与 TensorFlow 不同，PyTorch 使用动态类型图，这意味着执行图是在运行中创建的。它允许我们随时修改和检查图的内部结构；

• 除了用户友好的高级 API 之外，PyTorch 还包括精心构建的低级 API，允许对机器学习模型进行越来越多的控制。我们可以在训练期间对模型的前向和后向传递进行检查和修改输出。这被证明对于梯度裁剪和神经风格迁移非常有效；

• PyTorch 允许用户扩展代码，可以轻松添加新的损失函数和用户定义的层。PyTorch 的 Autograd 模块实现了深度学习算法中的反向传播求导数，在 Tensor 类上的所有操作， Autograd 都能自动提供微分，简化了手动计算导数的复杂过程；

• PyTorch 对数据并行和 GPU 的使用具有广泛的支持；

• PyTorch 比 TensorFlow 更 Python 化。PyTorch 非常适合 Python 生态系统，它允许使用 Python 类调试器工具来调试 PyTorch 代码。

• 正如官方网站所描述的那样，JAX 能够执行 Python+NumPy 程序的可组合转换：向量化、JIT 到 GPU/TPU 等等；

• 与 PyTorch 相比，JAX 最重要的方面是如何计算梯度。在 Torch 中，图是在前向传递期间创建的，梯度在后向传递期间计算， 另一方面，在 JAX 中，计算表示为函数。在函数上使用 grad() 返回一个梯度函数，该函数直接计算给定输入的函数梯度；

• JAX 是一个 autograd 工具，不建议单独使用。有各种基于 JAX 的机器学习库，其中值得注意的是 ObJax、Flax 和 Elegy。由于它们都使用相同的核心并且接口只是 JAX 库的 wrapper，因此可以将它们放在同一个 bracket 下；

• Flax 最初是在 PyTorch 生态系统下开发的，更注重使用的灵活性。另一方面，Elegy 受 Keras 启发。ObJAX 主要是为以研究为导向的目的而设计的，它更注重简单性和可理解性。 

© THE END 

投稿或寻求报道：[email protected]