機器之心報道
參與：思
感覺深度學習建模只不過調庫與堆疊層級？你需要谷歌大腦維護的這條路徑 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 腳本庫或者 Jupyter Notebook 的基礎，也可以作為命令行工具執行。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，然後如下運行就行了：
<code>python -m trax.trainer --config_file=$PWD/trax/configs/mlp_mnist.gin/<code>
如果你覺得上面的訓練太簡單，也可以在 ImageNet64 上訓練一下 Reformer：
<code>python -m trax.trainer --config_file=$PWD/trax/configs/reformer_imagenet64.gin/<code>
最後，這個項目最重要的還是它的實現代碼，我們並不是因為可以直接運行而使用它。相反，我們是因為它的代碼直觀簡潔，能幫助我們一步步更深刻地理解模型而使用它。
閱讀更多 機器之心Pro 的文章
