限时领|《AI 百问百答》专栏课+实体书(包邮)! 了解详情
写点什么

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


2018-12-12 07:002654
用户头像
陈思 InfoQ编辑

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

关注

评论 1 条评论

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

怎么清空.NET数据库连接池

喵叔

11月日更

npm必知必会点

废材壶

大前端 npm Node

赞!一篇博客讲解清楚 Python queue模块,作为Python爬虫预备知识,用它解决采集队列问题

梦想橡皮擦

11月日更

这一次,Google 终于对 Web 自动化下手了!

星安果

chrome 自动化

Android C++系列:JNI操作Bitmap

轻口味

c++ android jni 11月日更

如何评价一个开源项目(一)--活跃度

腾源会

开源

腾讯自研分布式远程Shuffle服务Firestorm正式开源

腾源会

大数据 开源 腾讯

我在 IBM 从事开源工作的十一年

腾源会

开源

Prometeus 2.31.0 新特性

耳东@Erdong

release Prometheus 11月日更

腾讯云原生开源生态专场召开,洞察开源云原生技术发展趋势和商业化路径

腾源会

腾讯云 开源 云原生

【高并发】通过源码深度解析ThreadPoolExecutor类是如何保证线程池正确运行的

冰河

Java 并发编程 多线程 高并发 异步编程

面试官:讲讲雪花算法,越详细越好

秦怀杂货店

分布式 雪花算法

SuperEdge 和 FabEdge 联合在边缘 K8s 集群支持原生 Service 云边互访和 PodIP 直通

腾源会

开源 边缘计算 superedge

腾讯开源全景图再刷新:社区贡献领跑国内企业,获超过38万开发者关注

腾源会

开源 腾讯

范学雷的专栏《深入剖析 Java 新特性》

IT蜗壳-Tango

11月日更

CNCF 沙箱再添“新将”!云原生边缘容器开源项目 SuperEdge 正式入选

腾源会

开源 容器 云原生 cncf

数据库连接池Demo(1)单线程初步

Java 数据库 连接池

Ubuntu系统下《汇编语言》环境配置

codists

汇编语言

进击的Java(七)

ES_her0

11月日更

消息队列表设计

Rabbit

Serverless 架构模式及演进

阿里巴巴云原生

阿里云 Serverless 云原生 架构模式

你不知道的开源分布式存储系统 Alluxio 源码完整解析(下篇)

腾源会

大数据 开源

在华为云专属月,找到开启互联网第二增长曲线的一把钥匙

脑极体

多模态内容理解算法框架项目 Lichee 正式开源,为微服务开源社区贡献力量

腾源会

开源

Github webhooks 自动部署博客文章,使用总结【含视频】

小傅哥

GitHub 小傅哥 WEBHOOKS 自动部署 通知回调

一文告诉你 K8s PR (Pull Request) 怎样才能被 merge?

腾源会

k8s

Golang Gin 框架入门介绍(二)

liuzhen007

11月日更

架构训练营 模块三 作业

dog_brother

「架构实战营」

腾讯发布 K8s 多集群管理开源项目 Clusternet

腾源会

开源 K8s 多集群管理 Clusternet

[ CloudWeGo 微服务实践 - 08 ] Nacos 服务发现扩展 (2)

baiyutang

golang 微服务 11月日更

干货分享:细说双 11 直播背后的压测保障技术

阿里巴巴云原生

阿里云 云原生 性能测试 PTS

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