【QCon】精华内容上线85%,全面覆盖“人工智能+”的典型案例!>>> 了解详情
写点什么

TensorFlow 官方简化版!谷歌开源机器学习库 JAX

整理编译自Reddit

  • 2018-12-12
  • 本文字数:2353 字

    阅读完需:约 8 分钟

TensorFlow官方简化版!谷歌开源机器学习库JAX

AI 前线导读:什么?TensorFlow 有了替代品?什么?竟然还是谷歌自己做出来的?先别慌,从各种意义上来说,这个所谓的“替代品”其实是 TensorFlow 的一个简化库,名为 JAX,结合 Autograd 和 XLA,可以支持部分 TensorFlow 的功能,但是比 TensorFlow 更加简洁易用。虽然还不至于替代 TensorFlow,但已经有 Reddit 网友对 JAX 寄予厚望,并表示“早就期待能有一个可以直接调用 Numpy API 接口的库了!”,“希望它可以取代 TensorFlow!”。


更多干货内容请关注微信公众号“AI 前线”(ID:ai-front)


JAX 结合了 Autograd 和 XLA,是专为高性能机器学习研究打造的产品。



有了新版本的 Autograd,JAX 能够自动对 Python 和 NumPy 的自带函数求导,支持循环、分支、递归、闭包函数求导,而且可以求三阶导数。它支持自动模式反向求导(也就是反向传播)和正向求导,且二者可以任意组合成任何顺序。


JAX 的创新之处在于,它基于 XLA 在 GPU 和 TPU 上编译和运行 NumPy 程序。默认情况下,编译是在底层进行的,库调用能够及时编译和执行。但是 JAX 还允许使用单一函数 API jit将自己的 Python 函数及时编译成经过 XLA 优化的内核。编译和自动求导可以任意组合,因此可以在不脱离 Python 环境的情况下实现复杂算法并获得最优性能。


JAX 最初由 Matt Johnson、Roy Frostig、Dougal Maclaurin 和 Chris Leary 发起,他们均任职于谷歌大脑团队。在 GitHub 的说明文档中,作者明确表示:JAX 目前还只是一个研究项目,不是谷歌的官方产品,因此可能会有一些 bug。从作者的 GitHub 简介来看,这应该是谷歌大脑正在尝试的新项目,在同一个 GitHub 目录下的开源项目还包括 8 月份在业内引起热议的强化学习框架 Dopamine。


以下是 JAX 的简单使用示例。



GitHub 项目传送门:https://github.com/google/JAX


有关具体的安装和简单的入门指导大家可以在 GitHub 中自行查看,在此不做过多赘述。

JAX 库的实现原理

机器学习中的编程是关于函数的表达和转换。转换包括自动微分、加速器编译和自动批处理。像 Python 这样的高级语言非常适合表达函数,但是通常使用者只能应用它们。我们无法访问它们的内部结构,因此无法执行转换。


JAX 可以用于专门化高级 Python+NumPy 函数,并将其转换为可转换的表示形式,然后再提升为 Python 函数。



JAX 通过跟踪专门处理 Python 函数。跟踪一个函数意味着:监视应用于其输入,以产生其输出的所有基本操作,并在有向无环图(DAG)中记录这些操作及其之间的数据流。为了执行跟踪,JAX 包装了基本的操作,就像基本的数字内核一样,这样一来,当调用它们时,它们就会将自己添加到执行的操作列表以及输入和输出中。为了跟踪这些原语之间的数据流,跟踪的值被包装在 Tracer 类的实例中。


当 Python 函数被提供给 grad 或 jit 时,它被包装起来以便跟踪并返回。当调用包装的函数时,我们将提供的具体参数抽象到 AbstractValue 类的实例中,将它们框起来用于跟踪跟踪器类的实例,并对它们调用函数。


抽象参数表示一组可能的值,而不是特定的值:例如,jit 将 ndarray 参数抽象为抽象值,这些值表示具有相同形状和数据类型的所有 ndarray。相反,grad 抽象 ndarray 参数来表示底层值的无穷小邻域。通过在这些抽象值上跟踪 Python 函数,我们确保它足够专门化,以便转换是可处理的,并且它仍然足够通用,以便转换后的结果是有用的,并且可能是可重用的。然后将这些转换后的函数提升回 Python 可调用函数,这样就可以根据需要跟踪并再次转换它们。


JAX 跟踪的基本函数大多与 XLA HLO 1:1 对应,并在 lax.py 中定义。这种 1:1 的对应关系使得到 XLA 的大多数转换基本上都很简单,并且确保我们只有一小组原语来覆盖其他转换,比如自动微分。 jax.numpy 层是用纯 Python 编写的,它只是用 LAX 函数(以及我们已经编写的其他 numpy 函数)表示 numpy 函数。这使得 jax.numpy 易于延展。


当你使用 jax.numpy 时,底层 LAX 原语是在后台进行 jit 编译的,允许你在加速器上执行每个原语操作的同时编写不受限制的 Python+ numpy 代码。


但是 JAX 可以做更多的事情:你可以在越来越大的函数上使用 jit 来进行端到端编译和优化,而不仅仅是编译和调度到一组固定的单个原语。例如,可以编译整个网络,或者编译整个梯度计算和优化器更新步骤,而不仅仅是编译和调度卷积运算。


折衷之处是,jit 函数必须满足一些额外的专门化需求:因为我们希望编译专门针对形状和数据类型的跟踪,但不是专门针对具体值的跟踪,所以 jit 装饰器下的 Python 代码必须适用于抽象值。如果我们尝试在一个抽象的 x 上求 x >0 的值,结果是一个抽象的值,表示集合{True, False},所以 Python 分支就像 if x > 0 会引起报错。


有关使用 jit 的更多要求,请参见:https://github.com/google/jax#whats-supported


好消息是,jit 是可选的:JAX 库在后台对单个操作和函数使用 jit,允许编写不受限制的 Python+Numpy,同时仍然使用硬件加速器。但是,当你希望最大化性能时,通常可以在自己的代码中使用 jit 编译和端到端优化更大的函数。

后续计划

目前项目小组还将对以下几项做更多尝试和更新:


  1. 完善说明文档

  2. 支持 Cloud TPU

  3. 支持多 GPU 和多 TPU

  4. 支持完整的 NumPy 功能和部分 SciPy 功能

  5. 全面支持 vmap

  6. 加速

  7. 降低 XLA 函数调度开销

  8. 线性代数例程(CPU 上的 MKL 和 GPU 上的 MAGMA)

  9. 高效自动微分原语condwhile


有关 JAX 库的介绍大致如此,如果你在尝试了 JAX 之后有一些较好的使用心得,欢迎随时向我们投稿,AI 前线十分愿意将你的经验传播给更多开发者。


再次附上 GitHub 链接:https://github.com/google/jax


相关资源:


JAX 论文链接:https://www.sysml.cc/doc/146.pdf

会议推荐

AICon


公众号推荐:

2024 年 1 月,InfoQ 研究中心重磅发布《大语言模型综合能力测评报告 2024》,揭示了 10 个大模型在语义理解、文学创作、知识问答等领域的卓越表现。ChatGPT-4、文心一言等领先模型在编程、逻辑推理等方面展现出惊人的进步,预示着大模型将在 2024 年迎来更广泛的应用和创新。关注公众号「AI 前线」,回复「大模型报告」免费获取电子版研究报告。

AI 前线公众号
2018-12-12 07:002252
用户头像
陈思 InfoQ编辑

发布了 576 篇内容, 共 260.5 次阅读, 收获喜欢 1293 次。

关注

评论 1 条评论

发布
暂无评论
发现更多内容

在线文本右边批量删除字符工具

入门小站

工具

vue指令-2

小恺

6月月更

【100个 Unity踩坑小知识点】| Unity 使用Quaternion.AngleAxis随机一个方向

呆呆敲代码的小Y

大家的 Hexo 博客都还好吗?

jrwng

Hexo

数据类型

Jason199

js 数据类型 6月月更

spring4.1.8初始化源码学习三部曲之一:AbstractApplicationContext构造方法

程序员欣宸

Java spring 6月月更

理解 Java 中的 NumberFormatException 异常

HoneyMoose

Junit 测试中如何对异常进行断言

HoneyMoose

如何使用物联网低代码平台进行模型管理?

AIRIOT

物联网 低代码开发

PHP开发者福音,支持CRUD代码生成且前后分离的tp6+Vue3后台管理系统开源啦!

妙码生花

php typescript Vue3 thinkphp Pinia

数据库每日一题---第4天:从不订购的客户

知心宝贝

数据库 程序员 前端 后端 6月月更

Django API 开发:博客系统的权限管理

宇宙之一粟

django 6月月更

NFT赛道或进入聚合时代,OKALEIDO成BNB Chain上的首个NFT聚合平台

股市老人

【100个 Unity踩坑小知识点】| Unity中的 碰撞盒检测 Physics.OverlapBox、OverlapCapsule、OverlapSphere

呆呆敲代码的小Y

【100个 Unity踩坑小知识点】| 在编辑器中绘制正方体虚线、球体虚线(Gizmos 辅助线框)

呆呆敲代码的小Y

运维服务体系构建

阿泽🧸

运维体系 6月月更

this和super的用法与区别

写代码两年半

继承 super javase this 6月月更

如何写好技术博客

卢卡多多

技术 博客 6月月更

AssertJ 的异常(Exception )断言

HoneyMoose

7 个 Flutter 开源项目,让你成为更好的 Flutter 开发者

坚果

6月月更

在线JADE转HTML工具

入门小站

工具

【100个 Unity实用技能】| 游戏中获取鼠标点击的坐标,并将游戏对象移动到鼠标的点击位置

呆呆敲代码的小Y

【100个 Unity踩坑小知识点】| Unity中的 Development build 详细解析

呆呆敲代码的小Y

【网络编程知识】什么是Socket?概念及原理分析

呆呆敲代码的小Y

【100个 Unity实用技能】| 游戏中使技能或装备跟随角色环绕,持续旋转

呆呆敲代码的小Y

【100个 Unity踩坑小知识点】| Unity控制物体持续指向某个方向

呆呆敲代码的小Y

见自身,见天地:华为,将ICT技术变成绿色未来的支点

脑极体

【100个 Unity踩坑小知识点】| Unity调用API ,动态获取Android权限,附带所有Android权限表格

呆呆敲代码的小Y

【愚公系列】2022年06月 面向对象设计原则(一)-单一职责原则

愚公搬代码

6月月更

软件设计本质

GalaxyCreater

架构 设计模式 架构师 系统 软件设计

Fabric.js 精简输出的JSON🎫

德育处主任

FabricJS 6月月更

TensorFlow官方简化版!谷歌开源机器学习库JAX_AI&大模型_InfoQ精选文章