谷歌JAX快速入门笔记详解和案例
附:推荐一本新书:《JAX可微分编程》,写的比较好,这里是微信公众号链接: 北大天才00后少年新作出版,这本书满足你对JAX的好奇心!
一. 什么是JAX?
JAX最初由谷歌大脑团队的 Matt Johnson、Roy Frostig、Dougal Maclaurin 和 Chris Leary 等人发起,借助 Autograd 的更新版本,并且结合了 XLA,可对 Python 程序与 NumPy 运算执行自动微分,支持循环、分支、递归、闭包函数求导,也可以求三阶导数;依赖于 XLA,JAX 可以在 GPU 和 TPU 上编译和运行 NumPy 程序;通过 grad,可以支持自动模式反向传播和正向传播,且二者可以任意组合成任何顺序。
JAX并非是一个深度学习的框架或者库,它的设计目标也并非是作为一个新的深度学习框架。
简单来说,JAX是一个包含可组合函数变换的数值计算库,只不过深度学习恰好是JAX能做的一项工作。
JAX处于函数变换(function transformations)和科学计算的交界处,所以也有能力训练神经网络模型,但不止于训练。
目前JAX在Github上已经斩获了超2万多颗star:
github地址: https://github.com/google/jax(截至目前,star数:20.3k)
官方文档: https://jax.readthedocs.io/en/latest/
JAX 是一个非常有前途的项目,并且用户一直在稳步增长。JAX 已经在深度学习、机器人 / 控制系统、贝叶斯方法和科学模拟等诸多领域得到了广泛应用。
二. 为什么应该使用JAX
JAX目前已经达到深度学习的最高水平。在当前开源的框架中,没有哪一个框架能在简洁、易用、速度
这3个方面有两个能同时超过JAX。
简洁
:JAX的设计追求最少的封装,尽量避免重复造轮子。设计遵循tensor→variable(autograd)→module
3个由低到高的层次,分别代表高维数组(张量),自动求导(变量)和神经网络(层/模块),而且这3个抽象直接连接紧密,可以同时进行修改和操作。而tensorflow充斥着graph、operation、tensor、layer等全新的概念。JAX源码只有 tensorflow 的十分之一左右,更少的抽象、更直观的设
2401_85111901: 支持一下,细节很到位!【我也写了一些相关领域的文章,希望能够得到博主的指导,共同进步!】
ak2111: 简明扼要地阐述了概念,大幅提升了文章的实用性和参考价值。整体而言,启发性与指导性并存,专业性和实用性俱佳。
ak2111: 本文以清晰易懂的语言深入解释了复杂概念,结合实用示例,逻辑严谨,观点清晰,引人入胜,是深思熟虑后的结果。结构合理,信息分层明确,让人不由得沉浸其中,受益匪浅。
ak2111: 使用设计模式是为了重用代码、让代码更容易被他人理解、保证代码可靠性。 毫无疑问,设计模式于己于他人于系统都是多赢的,设计模式使代码编制真正工程化,设计模式是软件工程的基石,如同大厦的一块块砖石一样。
2401_84166258: 文章内容通俗易懂,适合不同层次的读者。【我也写了一些相关领域的文章,希望能够得到博主的指导,共同进步!】