<引用类=" js_blockquote_wrap article-blockquote ">
首先值得注意的是俳句,这是一个面向Jax的深度学习代码库,它是由十四行诗作者——一个谷歌的神经网络库团队开发的。
为什么要使用俳句?这是因为其支持的是Jax, Jax在灵活性和性能上具有相当的优势。但是另一方面,Jax本身是函数式的,和面向对象的用户习惯有差别。因此,通过俳句,用户可以在Jax上进行面向对象开发了。
此外,俳句的API和编程模型都是基于十四行诗,因此使用过十四行诗的用户可以快速上手。项目作者也表示,十四行诗之于TensorFlow的提升就如同俳句之于Jax。
目前,俳句已公开了Alpha版,本已完全开源。项目作者欢迎使用者提出建议。
俳句主要分为两个模块:香港。模块和hk.transform。下文将会分别介绍。
香港。Python模块是对象,保存着到参数,其他模块和方法的参照(引用)。
香港。变换则负责将面向对象的模块转换为纯粹的函数式代码,然后让jax中的jax。jit, jax。校友,jax。pmap等进行处理,从而实现和Jax组件的兼容。
俳句能够做到很多机器学习需要完成的任务,相关功能和代码如下:
有时候,我们想要在训练的过程中保持某些内部参数的状态,在俳句上这也是非常容易实现的。
由于所有的代码都会被转换成Jax的函数,因此它们和jax.pmap。是完全兼容的。这说明,我们可以利用jax。pmap来进行分布式计算。
如下为进行数据分割的分布式加速代码,首先,我们先定义模型和训练步骤:
除了令人印象深刻的俳句外,DeepMind还开源了RLax——这是一个基于Jax的强化学习库。
相比俳句,RLax专门针对强化学习。项目作者认为,尽管强化学习中的算子和函数并不是完全的算法,但是,如果需要构建完全基于函数式的智能体,就需要特定的数学算子。
因此,函数式的Jax就成为了一个不错的选择。在Jax上进行一定的开发后,就可以有专用的强化学习库了.RLax目前的资料还较少,但项目已提供了一个示例代码:使用RLax进行q学习模型的搭建和训练。
代码如下,首先,使用俳句构建基本的强化学习模型:
