新星JAX :双挑TensorFlow和PyTorch!有望担纲Google主要科学计算...
新智元
共 2945字,需浏览 6分钟
· 2020-02-24
新智元报道
编辑:SHAO,元子
【新智元导读】JAX是机器学习框架领域的新生力量,它具有更快的高阶渐变,它建立在XLA之上,可以更快或在将来具有其他优势,具有其他有趣的转换和更好的TPU支持,甚至将来可能会成为Google的主要科学计算/NN库。「新智元急聘主笔、编辑、运营经理、客户经理,添加HR微信(Dr-wly)或扫描文末二维码了解详情。」
JAX:CPU、GPU和TPU上的Numpy
JAX官方文档是这样解释的:“JAX是CPU、GPU和TPU上的NumPy,具有出色的自动差异化功能,可用于高性能机器学习研究。”
就像文档上说的那样,最简单的JAX是加速器支持的numpy,它具有一些便利的功能,用于常见的机器学习操作。
从2006年开始,你就可以得到numpy精心设计的API,它具有像Tensorflow和PyTorch这样的现代机器学习工具的性能特征。
通过 jax.scipy,JAX还包括scipy项目的很大一部分。
尽管加速器支持的numpy + scipy版本已经非常有用,但JAX还有一些其他的妙招。首先让我们看看JAX对自动微分的支持。
JAX的前身:Autograd
Autograd是一个用于通过numpy和本机python代码高效计算梯度的库。Autograd也恰好是JAX的(很大程度上是字面意义)前身。尽管原始的autograd存储库已经不再重点开发 ,但是许多从事autograd工作的核心团队已全职转向于JAX项目。
就像autograd的名字一样,JAX允许通过简单地调用grad来获取python函数输出导数:
您还可以通过本机python控制结构进行微分 —— 不需要与tf.cond较劲。
JAX支持求解高阶导数,grad函数可以任意嵌套使用
默认情况下,grad为您提供了反向模式梯度,这是计算梯度的最常见模式,它依赖于缓存激活来使反向传递高效。反向模式差分通常是计算参数更新的最有效方法。但是,尤其是在实施依赖于高阶导数的优化方法时,它并不总是最佳选择。JAX通过jacfwd和jacrev对反向和正向模式自动微分提供优异的支持:
除了grad、jacfwd和jacrev之外,JAX还提供了计算函数的线性近似值、定义自定义梯度操作等实用程序,作为其自动微分支持的一部分。
XLA:将JAX转化为加速器支持操作的中坚力量
XLA(加速线性代数)是一个线性代数代码的特定领域编译器,它是允许JAX将python和numpy表达式,转化为加速器支持的操作的中坚力量。
除了允许JAX将python + numpy代码转换为可以在加速器上运行的操作(如我们在第一个示例中看到的那样)之外,XLA还允许JAX将几个操作融合到一起。它在计算图中寻找节点的簇,以减少计算或变量的中间存储。Tensorflow关于XLA的文档中,使用下面的例子来解释会从XLA编译中受益的实例。
在没有XLA的情况下运行,这将作为3个独立的内核运行——乘法、加法和加法归约。在有XLA运行的时候,这变成了一个负责这三者的单一内核,因为不需要存储中间变量,从而节省了时间和内存。
启用对此操作重写的支持与使用就像用@jax.jit来修饰一个函数一样简单:
像所有其他JAX函数一样,jax.jit是完全可组合的:
尽管Autograd和XLA构成了JAX库的核心,但是还有另外两个JAX函数脱颖而出。您可以使用jax.vmap和jax.pmap进行矢量化和基于SPMD的(单程序多数据)并行。
为了说明vmap的好处,我们将返回简单密集层的示例,该层在向量x表示的单个示例上运行。
我们已经将隐藏层编写为接受单个向量输入,但是实际上,我们几乎总是将输入分批处理以利用向量化计算。使用JAX,您可以使用任何接受单个输入并允许其接受一批输入的函数jax.vmap:
这其中的美妙之处在于,它意味着你或多或少地忽略了模型函数中的批处理维度,并且在你构建模型的时候,在你的头脑中总是少了一个张量维度。
如果您有多个应该全部矢量化的输入,或者要沿除轴0以外的其他轴矢量化,则可以使用in_axes参数指定此输入。
JAX的SPMD并行处理实用程序遵循非常相似的API。如果您有一台4-gpu的计算机,并且有一批4个示例,则可以使用pmap每个设备运行一个示例。
像往常一样,你可以随心所欲地编写函数。
相比TensorFlow和PyTorch,JAX的优势是什么?
我们使用的工具对我们探索的研究领域有着不相称的影响。无论是有意识的,还是在潜意识里,我们把想法限制在我们知道如何有效实施的想法空间里。正如罗曼·林(Roman Ring)在他探索下一代机器学习工具的博客文章中指出的那样,AlexNet主要是一项软件工程成就,它使得几十年来的好的机器学习想法能够得到正确的测试。
JAX向前迈出了重要的一步,不是因为它比现有的机器学习框架具有更简洁的API,或者因为它比Tensorflow和PyTorch在做它们被设计的事情上做得更好,而是因为它允许我们更容易地尝试以前可能的更广阔的思想空间。
如果您深入研究并开始将JAX用于自己的项目,你可能会对JAX在表面上做得如此之少而感到沮丧。需要手工编写训练循环,管理参数需要自定义代码。每当你想要一个新的随机值时,你甚至必须生成你自己的随机PRNG密钥。但在某种程度上,这也是JAX最大的优势。
它不会把你看不到的细节藏在窗帘后面。内部结构被广泛地记录下来,很明显,JAX关心的是让其他开发者做出贡献。JAX对你打算如何使用它做了很少的假设,这样做给了你在其他框架中做不到的灵活性。
每当您将一个较低的API封装到一个较高的抽象层时,您就要对最终用户可能拥有的使用空间做出假设。当您心中有一个非常有针对性的应用程序时,这就形成了非常简洁的应用程序接口,允许您用最少的配置获得想要的结果。特别是最近在TF2.0中强调了Keras和更高级的APIs,编写Tensorflow感觉有点像使用3D打印机一样简单,只要你想要一个适合打印表面的塑料物体,它就会像被施了魔法一样工作。
和JAX一起工作就像被允许进入一个功能完善的机械车间。是的,如果你不小心的话,你周围的一切东西,都有可能让你被砍掉一根手指或者造成严重的身体伤害。但是,自由地实施和探索那些可能行得通的“无处不在”的想法,使得利用JAX变得值得。所以请戴上你隐喻性的护目镜,开始使用JAX建造一些奇怪的东西。
尽管JAX的生态系统仍然相当分散,但是确实存在一些在JAX之上构建的框架,这些框架在核心应用编程接口之上提供了一些简单的抽象。特别值得注意的是一下几个:
Flax:为灵活性而设计的功能框架
Trax: [Tensor2Tensor]的精神传承者,拥有TF和JAX的后台支持。
Stax:一个神经网络库,是建成的JAX实验模块的一部分
JAX生态系统中的所有库似乎都以“ax”结尾。
评论
如何计算数据中心的冷却需求?
今日分享 【导读】数据中心的冷却要求受多种因素影响,包括设备的热量输出、占地面积、设施设计和电气系统功率额定值等等……众所周知,环境因素会严重影响数据中心设备。过多的热量积聚会损坏服务器,可能导致其自动关闭。经常在高于可接受的温度下运行服务器会缩短其使用
数据中心运维管理
0
Open-Sora全面开源升级:支持16s视频生成和720p分辨率
机器之心发布 机器之心编辑部Open-Sora 在开源社区悄悄更新了,现在单镜头支持长达16秒的视频生成,分辨率最高可达720p,并且可以处理任何宽高比的文本到图像、文本到视频、图像到视频、视频到视频和无限长视频的生成需求。我们来试试效果。生成个横屏圣诞雪景,发b站再生成个竖屏,发抖音还能
机器学习算法与Python实战
0
聊一聊我最关注的9个CV、SLAM、自动驾驶和AI圈子!
随着计算机视觉(2D/3D)、SLAM、自动驾驶、AI技术的快速迭代更新,可落地的技术也成为人们争先学习的重点。这使得从业者对于最前沿技术的获取能力变得至关重要。微信公众号便是一个非常有效的前沿信息分享平台。这里给大家推荐9个最常打开的计算机视觉、自动驾驶、SLAM、机器学习和AI方向的优质公众号平
3D视觉工坊
0
文本嵌入、文本分类和语义搜索
在实践中使用大型语言模型(LLM)中,RAG 的一个关键部分是使用文本嵌入从知识库中自动检索相关信息。在这里,我将更深入地讨论文本嵌入,并分享两个简单(但功能强大)的应用:文本分类和语义搜索。ChatGPT 吸引了全世界对人工智能及其潜力的想象。ChatGPT 的聊天界面是这一影响的关键因素,它使人
大邓和他的Python
0
学习开放日:开放复杂科学、AI+X 海量学习资源!
Datawhale干货 学习开放日:4月27-28日1. 什么是学习开放日?以AI为代表的技术突飞猛进,人类知识森林快速扩张,仅凭一人之力不仅难以覆盖,更是难以串联知识线索。唯有像蚂蚁探索最优路径一样,我们才能在信息爆炸的知识森林中探索出更好的方向!因此,今年集智斑图联合国内最
Datawhale
1
一站式解决方案:基于 Arthas 实现服务发现和权限控制
来源:juejin.cn/post/7281849496983994383👉 欢迎加入小哈的星球 ,你将获得: 专属的项目实战 / Java 学习路线 / 一对一提问 / 学习打卡 / 赠书福利全栈前后端分离博客项目 2.0 版本完结啦, 演示链接
小哈学Java
0
Go早期是如何在Google内部发展起来的
2007年Go诞生于Google,2009年Google正式对外宣布了Go语言的开源!时至今日,距离Go开源已经过去了近15个年头了[1]!Go在Google公司内部究竟是怎样的一个状态呢?前Google员工Yves Junqueira近期撰文从其个人所见所闻谈了Go在Google的历程[2]!这里
GoCN
0
互联网晚报 | 大麦网已退款凤凰传奇演唱会“柱子票”;钟薛高再成被执行人;iPhone 16或取消实体音量键和电源键
大麦网回应凤凰传奇演唱会买到“柱子票”:已退票退款据报道,凤凰传奇2024巡回演唱会常州站演出结束的第二天,有网友称自己在大麦网买到“柱子票”,因为观看效果不佳,要求退款被拒。23日,记者从涉事网友处了解到,大麦方面给出了退款建议,但被其拒绝,“我希望平台退款加赔偿,并重视屡次出现的‘柱子票’问题。
产品刘
0