谷歌开源机器学习库JAX
极客时间编辑部
讲述:丁婵大小:3.37M时长:02:27
近日,谷歌开源了机器学习库 JAX,号称要替代 TensorFlow,各种意义上来说,这个所谓的“替代品”其实是 TensorFlow 的一个简化库,结合 Autograd 和 XLA,可以支持部分 TensorFlow 的功能,但是比 TensorFlow 更加简洁易用。虽然还不至于替代 TensorFlow,但是已经有 Reddit 网友对 JAX 寄予厚望。
有了新版本的 Autograd,JAX 能够自动对 Python 和 NumPy 的自带函数求导,支持循环、分支、递归、闭包函数求导,而且可以求三阶导数。它支持自动模式反向求导(也就是反向传播)和正向求导,二者可以任意组合成任何顺序。
据介绍,JAX 的创新之处在于,它基于 XLA 在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下,编译是在底层进行的,库调用能够及时编译和执行。但是 JAX 还允许使用单一函数 API jit,将自己的 Python 函数及时编译成经过 XLA 优化的内核。编译和自动求导可以任意组合,因此,可以在不脱离 Python 环境的情况下,实现复杂算法并获得最优性能。
在 GitHub 的说明文档中,作者明确表示,JAX 目前还只是一个研究项目,不是谷歌的官方产品,因此可能会有一些 bug。从作者的 GitHub 简介来看,这应该是谷歌大脑正在尝试的新项目,另外,在同一个 GitHub 目录下的开源项目,还包括 8 月份在业内引起热议的强化学习框架 Dopamine。
事实上,机器学习中的编程是关于函数的表达和转换。转换包括自动微分、加速器编译和自动批处理。像 Python 这样的高级语言非常适合表达函数,但是通常使用者只能应用它们,无法访问它们的内部结构,因此无法执行转换。
JAX 可以用于专门化高级 Python+NumPy 函数,并将其转换为可转换的表示形式,然后再提升为 Python 函数。
公开
同步至部落
取消
完成
0/2000
荧光笔
直线
曲线
笔记
复制
AI
- 深入了解
- 翻译
- 解释
- 总结
该免费文章来自《极客视点》,如需阅读全部文章,
请先领取课程
请先领取课程
免费领取
© 版权归极客邦科技所有,未经许可不得传播售卖。 页面已增加防盗追踪,如有侵权极客邦将依法追究其法律责任。
登录 后留言
全部留言(1)
- 最新
- 精选
- caohuantensorflow 还没搞好,又来了 JAX 技术变化很快,跟不上,没能更快,可以往底层钻研
收起评论