谷歌大脑开源Trax代码库,你的深度学习进阶路径

感觉深度学习建模只不过调库与堆叠层级?你需要谷歌大脑维护的这条路径 Trax,从头实现深度学习模型。

从最开始介绍卷积、循环神经网络原理,到后来展示各种最前沿的算法与论文,机器之心与读者共同探索着机器学习。我们会发现,现在读者对那些著名的深度学习模型已经非常熟悉了,经常也会推导或复现它们。

而对于最前沿的一些实现,包括 Transformer 或其它强化学习,我们通常都需要看原作者开源的代码,或者阅读大厂的复现。出于速度等方面的考虑,这些实现通常会显得比较「隐晦」,理解起来不是那么直白。这个时候,你就需要谷歌大脑维护的 Trax,它是 ML 开发者进阶高级 DL 模型的路径。

Trax 是一个开源项目,它的目的在于帮助我们挖掘并理解高一阶的深度学习模型。谷歌大脑表示,该项目希望 Trax 代码做到非常整洁与直观,并同时令 Reformer 这类高阶深度学习达到最好的效果。

项目地址:https://github.com/google/trax

什么是 Trax

简单来说,Trax 就是一个代码库,它有点类似于一个极简的深度学习框架。只不过 Trax 关注什么样的代码能让读者更好地理解模型,而不只是关注加速与优化。

Trax 代码及其组织方式希望让我们从头理解深度学习,而不只是简单地调库。整个项目从最基础的数学部分开始,然后向上依次构建层级运算、模型运算,以及有监督与强化学习训练任务。

因为是进阶深度学习高级建模,Trax 还囊括了最前沿的研究结果,例如在 ICLR 2020 上做演讲报告的 Reformer。如下展示的是该项目的代码文件结构:

如果要从头理解并进阶深度学习,那么 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,然后如下运行就行了:

python -m trax.trainer --config_file=$PWD/trax/configs/mlp_mnist.gin

如果你觉得上面的训练太简单,也可以在 ImageNet64 上训练一下 Reformer:

python -m trax.trainer --config_file=$PWD/trax/configs/reformer_imagenet64.gin

最后,这个项目最重要的还是它的实现代码,我们并不是因为可以直接运行而使用它。相反,我们是因为它的代码直观简洁,能帮助我们一步步更深刻地理解模型而使用它。

我来评几句
登录后评论

已发表评论数()

相关站点

+订阅
热门文章