DeepMind发布神经网络、强化学习库,网友:推动JAX发展

作者: 十三 2020-02-21 15:33:44

本文经AI新媒体量子位(公众号ID:QbitAI)授权转载,转载请联系出处。

 DeepMind今日发布了Haiku和RLax两个库,都是基于JAX。

JAX由谷歌提出,是TensorFlow的简化库。结合了针对线性代数的编译器XLA,和自动区分本地 Python 和 Numpy 代码的库Autograd,在高性能的机器学习研究中使用。

而此次发布的两个库,分别针对神经网络和强化学习,大幅简化了JAX的使用。

Haiku是基于JAX的神经网络库,允许用户使用熟悉的面向对象程序设计模型,可完全访问 JAX 的纯函数变换。

RLax是JAX顶层的库,它提供了用于实现增强学习代理的有用构件。

有意思的是,Reddit网友惊奇的发现Haiku这个库的名字,竟然不以“ax”结尾。

DeepMind发布神经网络、强化学习库,网友:推动JAX发展

当然,也有网友对这两个库表示了肯定:

毫无疑问,对JAX起到了推动作用。

DeepMind发布神经网络、强化学习库,网友:推动JAX发展

那么,我们就来看下Haiku和RLex的庐山真面目吧。

Haiku

Haiku是JAX的神经网络库,它允许用户使用熟悉的面向对象编程模型,同时允许完全访问JAX的纯函数转换。

它提供了两个核心工具:模块抽象hk.Module,和一个简单的函数转换hk.transform。

hk.Module是Python对象,包含对其自身参数、其他模块和对用户输入应用函数方法的引用。

hk.transform允许完全访问JAX的纯函数转换。

其实,在JAX中有许多神经网络库,那么Haiku有什么特别之处呢?有5点。

1、Haiku已经由DeepMind的研究人员进行了大规模测试

DeepMind相对容易地在Haiku和JAX中复制了许多实验。其中包括图像和语言处理的大规模结果、生成模型和强化学习。

2、Haiku是一个库,而不是一个框架

它的设计是为了简化一些具体的事情,包括管理模型参数和其他模型状态。可以与其他库一起编写,并与JAX的其他部分一起工作。

3、Haiku并不是另起炉灶

它建立在Sonnet的编程模型和API之上,Sonnet是DeepMind几乎普遍采用的神经网络库。它保留了Sonnet用于状态管理的基于模块的编程模型,同时保留了对JAX函数转换的访问。

4、过渡到Haiku是比较容易的

通过精心的设计,从TensorFlow和Sonnet,过渡到JAX和Haiku是比较容易的。除了新的函数(如hk.transform),Haiku的目的是Sonnet 2的API。

5、Haiku简化了JAX

它提供了一个处理随机数的简单模型。在转换后的函数中,hk.next_rng_key()返回一个唯一的rng键。

那么,该如何安装Haiku呢?

Haiku是用纯Python编写的,但是通过JAX依赖于c++代码。

首先,按照下方链接中的说明,安装带有相关加速器支持的JAX。

https://github.com/google/jax#installation

然后,只需要一句简单的pip命令就可以完成安装。

  1. $ pip install git+https://github.com/deepmind/haiku 

接下来,是一个神经网络和损失函数的例子。

  1. import haiku as hk 
  2.  
  3. import jax.numpy as jnp 
  4.  
  5. def softmax_cross_entropy(logits, labels): 
  6.  
  7.   one_hot = hk.one_hot(labels, logits.shape[-1]) 
  8.  
  9.   return -jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1
  10.  
  11. def loss_fn(images, labels): 
  12.  
  13.   model = hk.Sequential([ 
  14.  
  15.       hk.Linear(1000), 
  16.  
  17.       jax.nn.relu, 
  18.  
  19.       hk.Linear(100), 
  20.  
  21.       jax.nn.relu, 
  22.  
  23.       hk.Linear(10), 
  24.  
  25.   ]) 
  26.  
  27.   logits = model(images) 
  28.  
  29.   return jnp.mean(softmax_cross_entropy(logits, labels)) 
  30.  
  31. loss_obj = hk.transform(loss_fn) 

RLax

RLax是JAX顶层的库,它提供了用于实现增强学习代理的有用构件。

它所提供的操作和函数不是完整的算法,而是强化学习特定数学操作的实现。

RLax的安装也非常简单,一个pip命令就可以搞定。

  1. pip install git+git://github.com/deepmind/rlax.git 

使用JAX的jax.jit函数,所有的RLax代码可以不同的硬件上编译。

RLax需要注意的是它的命名规则。

许多函数在连续的时间步长中考虑策略、操作、奖励和值,以便计算它们的输出。在这种情况下,后缀_t和tm1通常是为了说明每个输入是在哪个步骤上生成的,例如:

q_tm1:转换的源状态中的操作值。

a_tm1:在源状态下选择的操作。

r_t:在目标状态下收集的结果奖励。

q_t:目标状态下的操作值。

Haiku和RLax都已在GitHub上开源,有兴趣的读者可从“传送门”的链接访问。

传送门

Haiku:

https://github.com/deepmind/haiku

RLax:

https://github.com/deepmind/rlax

人工智能 机器学习 技术
上一篇:谷歌透露:正在内部尝试用AI开发计算机芯片 下一篇:意料之外 情理之中:解读Gartner 2020年数据科学和机器学习平台魔力象限
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

12个场景应用,百余种算法,AI是如何攻占经济学的?

在虚拟世界中模拟现实经济状况,想法设计更好的制度只是AI和经济学结合方式之一。其实深度强化学习在面临风险参数和不确定性不断增加的现实经济问题时,也可以提供更好的性能和更高的精度。

蒋宝尚 ·  1天前
你在打王者农药,有人却用iPhone来训练神经网络

在 iOS 设备上也可以直接训练 LeNet 卷积神经网络,而且性能一点也不差,iPhone 和 iPad 也能化为实实在在的生产力。

佚名 ·  1天前
未来20年,全自动驾驶和互联汽车会出现吗?

在这篇博文中,重点介绍了汽车技术如何使交通更安全、更智能、更有趣。

Huibert Verhoeven ·  1天前
AI、机器学习和深度学习是 OEM 的主要市场

人工智能(AI)正在迅速改变全球行业参与者的经营方式。人工智能(AI)正在迅速改变全球行业参与者的经营方式。随着人工智能在商业和商业领域的广泛应用,我们看到了从更智能的产品到专注于聚焦客户服务的一切演变。

佚名 ·  1天前
和AI去码一样神奇?AI上色是黑科技还是逗你玩

在图片处理领域这块,AI 刷的存在感越来越多。早前笔者就介绍过 AI 无损放大图片、AI 去除马赛克、AI 自动给线稿上色之类的玩法,现在,又有人给笔者推荐了一个 AI 黑科技——黑白照片一键变彩色。

Aimo ·  1天前
AI与IoT:两种强大的技术将如何改变未来商业模式

人工智能和物联网正在重新定义企业过去的执行方式。在无人驾驶无人机和机器学习开始普及之前,詹姆斯·卡梅隆(James Cameron)于1984年通过他的梦想项目《终结者》(The Terminator)吸引了全世界的目光。

佚名 ·  1天前
为什么AI没能让人类失业?

我们听了这么久有关「AI取代设计师」、「AI取代写作者」、「AI取代画家」、「AI取代司机」……的论调。然而时至今日,我们发现实际上,目前还没有一个职业「真正」被取代。

佚名 ·  1天前
轻松构建 PyTorch 生成对抗网络(GAN)

生成对抗网络(GAN)是一种生成式机器学习模型,它被广泛应用于广告、游戏、娱乐、媒体、制药等行业,可以用来创造虚构的人物、场景,模拟人脸老化,图像风格变换,以及产生化学分子式等等。

佚名 ·  1天前
Copyright©2005-2020 51CTO.COM 版权所有 未经许可 请勿转载