首页 > AI资讯 > 最新资讯 > Github1.3万星,迅猛发展的JAX对比TensorFlow、PyTorch

Github1.3万星,迅猛发展的JAX对比TensorFlow、PyTorch

新火种    2023-09-20
JAX 是机器学习 (ML) 领域的新生力量,它有望使 ML 编程更加直观、结构化和简洁。在机器学习领域,大家可能对 TensorFlow 和 PyTorch 已经耳熟能详,但除了这两个框架,一些新生力量也不容小觑,它就是谷歌推出的 JAX。很对研究者对其寄予厚望,希望它可以取代 TensorFlow 等众多机器学习框架。JAX 最初由谷歌大脑团队的 Matt Johnson、Roy Frostig、Dougal Maclaurin 和 Chris Leary 等人发起。目前,JAX 在 GitHub 上已累积 13.7K 星。项目地址:https://github.com/google/jax迅速发展的 JAXJAX 的前身是 Autograd,其借助 Autograd 的更新版本,并且结合了 XLA,可对 Python 程序与 NumPy 运算执行自动微分,支持循环、分支、递归、闭包函数求导,也可以求三阶导数;依赖于 XLA,JAX 可以在 GPU 和 TPU 上编译和运行 NumPy 程序;通过 grad,可以支持自动模式反向传播和正向传播,且二者可以任意组合成任何顺序。开发 JAX 的出发点是什么?说到这,就不得不提 NumPy。NumPy 是 Python 中的一个基础数值运算库,被广泛使用。但是 numpy 不支持 GPU 或其他硬件加速器,也没有对反向传播的内置支持,此外,Python 本身的速度限制阻碍了 NumPy 使用,所以少有研究者在生产环境下直接用 numpy 训练或部署深度学习模型。在此情况下,出现了众多的深度学习框架,如 PyTorch、TensorFlow 等。但是 numpy 具有灵活、调试方便、API 稳定等独特的优势。而 JAX 的主要出发点就是将 numpy 的以上优势与硬件加速结合。目前,基于 JAX 已有很多优秀的开源项目,如谷歌的神经网络库团队开发了 Haiku,这是一个面向 Jax 的深度学习代码库,通过 Haiku,用户可以在 Jax 上进行面向对象开发;又比如 RLax,这是一个基于 Jax 的强化学习库,用户使用 RLax 就能进行 Q-learning 模型的搭建和训练;此外还包括基于 JAX 的深度学习库 JAXnet,该库一行代码就能定义计算图、可进行 GPU 加速。可以说,在过去几年中,JAX 掀起了深度学习研究的风暴,推动了科学研究迅速发展。JAX 的安装如何使用 JAX 呢?首先你需要在 Python 环境或 Google colab 中安装 JAX,使用 pip 进行安装:

$ pip install --upgrade jax jaxlib注意,上述安装方式只是支持在 CPU 上运行,如果你想在 GPU 执行程序,首先你需要有 CUDA、cuDNN ,然后运行以下命令(确保将 jaxlib 版本映射到 CUDA 版本):

$ pip install --upgrade jax jaxlib==0.1.61+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html现在将 JAX 与 Numpy 一起导入:

import jaximport jax.numpy as jnpimport numpy as npJAX 的一些特性使用 grad() 函数自动微分:这对深度学习应用非常有用,这样就可以很容易地运行反向传播,下面为一个简单的二次函数并在点 1.0 上求导的示例:

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.0jit(Just in time) :为了利用 XLA 的强大功能,必须将代码编译到 XLA 内核中。这就是 jit 发挥作用的地方。要使用 XLA 和 jit,用户可以使用 jit() 函数或 @jit 注释。

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 looppmap:自动将计算分配到所有当前设备,并处理它们之间的所有通信。JAX 通过 pmap 转换支持大规模的数据并行,从而将单个处理器无法处理的大数据进行处理。要检查可用设备,可以运行 jax.devices():

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)vmap:是一种函数转换,JAX 通过 vmap 变换提供了自动矢量化算法,大大简化了这种类型的计算,这使得研究人员在处理新算法时无需再去处理批量化的问题。示例如下:

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)TensorFlow vs PyTorch vs Jax在深度学习领域有几家巨头公司,他们所提出的框架被广大研究者使用。比如谷歌的 TensorFlow、Facebook 的 PyTorch、微软的 CNTK、亚马逊 AWS 的 MXnet 等。每种框架都有其优缺点,选择的时候需要根据自身需求进行选择。我们以 Python 中的 3 个主要深度学习框架——TensorFlow、PyTorch 和 Jax 为例进行比较。这些框架虽然不同,但有两个共同点:

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

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

那么它们的不同体现在哪些方面呢?如下表所示,为 TensorFlow、PyTorch、JAX 三个框架的比较。TensorFlowTensorFlow 由谷歌开发,最初版本可追溯到 2015 年开源的 TensorFlow0.1,之后发展稳定,拥有强大的用户群体,成为最受欢迎的深度学习框架。但是用户在使用时,也暴露了 TensorFlow 缺点,例如 API 稳定性不足、静态计算图编程复杂等缺陷。因此在 TensorFlow2.0 版本,谷歌将 Keras 纳入进来,成为 tf.keras。目前 TensorFlow 主要特点包括以下:

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

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

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

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

PyTorchPyTorch(Python-Torch) 是来自 Facebook 的机器学习库。用 TensorFlow 还是 PyTorch?在一年前,这个问题毫无争议,研究者大部分会选择 TensorFlow。但现在的情况大不一样了,使用 PyTorch 的研究者越来越多。PyTorch 的一些最重要的特性包括:

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

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

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

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

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

JAX JAX 是来自 Google 的一个相对较新的机器学习库。它更像是一个 autograd 库,可以区分原生的 python 和 NumPy 代码。JAX 的一些特性主要包括:

正如官方网站所描述的那样,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 主要是为以研究为导向的目的而设计的,它更注重简单性和可理解性。

Tags:
相关推荐
免责声明
本文所包含的观点仅代表作者个人看法,不代表新火种的观点。在新火种上获取的所有信息均不应被视为投资建议。新火种对本文可能提及或链接的任何项目不表示认可。 交易和投资涉及高风险,读者在采取与本文内容相关的任何行动之前,请务必进行充分的尽职调查。最终的决策应该基于您自己的独立判断。新火种不对因依赖本文观点而产生的任何金钱损失负任何责任。