江苏龙网

  1. 首页 > 人文 > >

科技实验室▲你的深度学习进阶路径，谷歌大脑开源Trax代码库

谷歌大脑开源Trax代码库 ， 你的深度学习进阶路径
科技实验室▲你的深度学习进阶路径，谷歌大脑开源Trax代码库
文章图片
机器之心Pro
1582785422机器之心报道
参与：思
感觉深度学习建模只不过调库与堆叠层级？你需要谷歌大脑维护的这条路径Trax ， 从头实现深度学习模型 。
从最开始介绍卷积、循环神经网络原理 ， 到后来展示各种最前沿的算法与论文 ， 机器之心与读者共同探索着机器学习 。 我们会发现 ， 现在读者对那些著名的深度学习模型已经非常熟悉了 ， 经常也会推导或复现它们 。
而对于最前沿的一些实现 ， 包括Transformer或其它强化学习 ， 我们通常都需要看原作者开源的代码 ， 或者阅读大厂的复现 。 出于速度等方面的考虑 ， 这些实现通常会显得比较「隐晦」 ， 理解起来不是那么直白 。 这个时候 ， 你就需要谷歌大脑维护的Trax ， 它是ML开发者进阶高级DL模型的路径 。
Trax是一个开源项目 ， 它的目的在于帮助我们挖掘并理解高一阶的深度学习模型 。 谷歌大脑表示 ， 该项目希望Trax代码做到非常整洁与直观 ， 并同时令Reformer这类高阶深度学习达到最好的效果 。
项目地址：https://github.com/google/trax
什么是Trax
简单来说 ， Trax就是一个代码库 ， 它有点类似于一个极简的深度学习框架 。 只不过Trax关注什么样的代码能让读者更好地理解模型 ， 而不只是关注加速与优化 。
Trax代码及其组织方式希望让我们从头理解深度学习 ， 而不只是简单地调库 。 整个项目从最基础的数学部分开始 ， 然后向上依次构建层级运算、模型运算 ， 以及有监督与强化学习训练任务 。
因为是进阶深度学习高级建模 ， Trax还囊括了最前沿的研究结果 ， 例如在ICLR2020上做演讲报告的Reformer 。 如下展示的是该项目的代码文件结构：
【科技实验室▲你的深度学习进阶路径，谷歌大脑开源Trax代码库】如果要从头理解并进阶深度学习 ， 那么Trax代码主要可以分为以下6部分：
math/：最基本的数学运算 ， 以及通过JAX和TensorFlow加速运算性能的方法 ， 尤其是在GPU/TPU上；layers/：搭建神经网络的所有层级构建块；models/：包含所有基础模型 ， 例如MLP、ResNet和Transformer ， 还包含一些前沿DL模型；optimizers/：包含深度学习所需要的最优化器；supervised/：包含执行监督学习的各种有用模块 ， 以及整体的训练工具；rl/：包含谷歌大脑在强化学习上的一些研究工作；每一个文件夹下都有对应的实现 ， 例如在Layers中 ， 所有神经网络层级都继承自最基础的Layer类 ， 实现这个类花了700行代码 。 而后新的层级在继承它后只要实现以下两个方法就行：
通过900行代码（包括Err处理） ， 基础的Layer类能完成其它所有处理 ， 包括初始化与调用等 。
使用Trax
我们可以将Trax作为Python脚本库或者JupyterNotebook的基础 ， 也可以作为命令行工具执行 。 Trax包含很多深度学习模型 ， 并且绑定了大量深度学习数据集 ， 包括Tensor2Tensor和TensorFlow采用的数据集 。 同时 ， 如果我们在CPU、GPU或TPU上运行这些模型 ， 也不需要改变 。
如果读者想要了解如何快速将Trax作为一个库来使用 ， 那么可以看看如下Colab上的入门示例 。 它介绍了如何生成样本数据 ， 并连接到Trax中的Transformer模型 。 在训练或推断时 ， 我们可以选择GPU ， 也可以选择8核心的免费TPU 。
入门简介地址：
https://colab.research.google.com/github/google/trax/blob/master/trax/intro.ipynb
如果要在命令行中使用Trax ， 那么带上参数就可以了 ， 例如模型类型、学习率等超参 。 谷歌大脑团队建议我们可以看看gin-config ， 例如训练一个最简单的MNIST分类模型 ， 可以看看mlp_mnist.gin ， 然后如下运行就行了：


推荐阅读


上一篇：『企业说』疫情期间降薪幅度最高达40%，优信又一高管离职、员工减薪

下一篇：『电气技术』英国谢菲尔德大学王涛特稿：非理想电网下双馈风电系统运行技术