TPU加AutoML：50美元快速訓練高效的ImageNet圖像分類網路

昨日，Jeff Dean 在推特上表示他們在 ImageNet 圖像分類上發布了新的 DAWNBench 結果，新型 AmoebaNet-D 網路在 TPU 上實現了最低的訓練時間和最少的訓練成本。在這一個基準測試上，基於進化策略的 DAWNBench 要比殘差網路效果更好，且訓練成本降低了一倍。機器之心簡要介紹了該基準測試和 AmoebaNet 網路架構，並提供了使用 TPU 進行 AmoebaNet 試驗的項目地址。

Jeff Dean 展示的兩個基準測試都是 DAWNBench 上的結果，DAWNBench 是斯坦福發布的一套基準測試，主要關注於端到端的深度學習訓練和推斷過程。雖然目前很多基準僅關注於模型精度，但 DAWNBench 格外關注計算時間和成本，它們是深度模型所需的關鍵資源。DAWNBench 提供了一套通用的深度學習工作負載參考集，因此能量化不同優化策略、模型架構、軟體框架、雲計算和硬體中的訓練時間、訓練成本、推斷延遲和推斷成本等重點指標。

DAWNBench 基準結果提交地址：https://github.com/stanford-futuredata/dawn-bench-entries

如下是各網路架構在圖像分類 Top-5 驗證準確率（ImageNet 數據集上）達到 93% 所需要的時間。Jeff Dean 所展示的 AmoebaNet-D 在同等的硬體條件下要比 50 層的殘差網路要快很多，它們所使用的框架也都展示在下圖。此外，DAWNBench 上所有的基準測試都需要提供測試源代碼。

下圖展示了不同模型 Top-5 準確率達到 93% 所需要的成本，因為 AmoebaNet-D 和 ResNet-50 使用相同的硬體，所以減少一半時間也就意味著成本減少了一半。

ResNet-50 已經是非常經典且成功的架構，即使我們將層級數增加到 100 至 150 層，網路的最終準確度也不會有一個質的提高。但新型的 AmoebaNet-D 大大減少了計算時間，這種網路基於進化策略能高效搜索神經網路架構並實現快速的訓練。谷歌上個月其實已經介紹了這種網路，詳細內容前查看：進化演算法 + AutoML，谷歌提出新型神經網路架構搜索方法。

AmoebaNet

在 ICML 2017 大會中展示的論文《Large-Scale Evolution of Image Classifiers》中，谷歌用簡單的構建模塊和常用的初始條件設置了一個進化過程。其主要思想是讓人「袖手旁觀」，讓進化演算法大規模構建網路架構。當時，從非常簡單的網路開始，該過程可以找到與手動設計模型性能相當的分類器。這個結果振奮人心，因為很多應用可能需要較少的用戶參與。例如，一些用戶可能需要更好的模型，但沒有足夠的時間成為機器學習專家。

接下來要考慮的問題自然就是手動設計和進化的組合能不能獲得比單獨使用一個方法更好的結果。因此，在近期論文《Regularized Evolution for Image Classifier Architecture Search》（即 AmoebaNet）中，谷歌通過提供複雜的構建模塊和較好的初始條件來參與進化過程。此外，谷歌還使用其新型 TPUv2 晶元來擴大計算規模。通過現代硬體、專家知識和進化過程的組合，谷歌獲得了在兩個流行的圖像分類基準 CIFAR-10 和 ImageNet 上的當前最優模型。

• AmoebaNet 論文：Regularized Evolution for Image Classifier Architecture Search

• 論文地址：https://arxiv.org/pdf/1802.01548.pdf

該論文中，群體中初始化的模型必須遵從 cell 外部堆疊（由專家設計）。這類似於 Zoph et al. (2017) 論文中的架構設計思想（如下圖）。儘管這些種子模型中的 cell 是隨機的，但是我們不再從簡單模型開始，這樣更易獲得高質量模型。

Zoph et al. (2017) 論文中提到的構建模塊。左側是完整神經網路的外部結構，自下而上地通過一串堆疊的重複 cell 解析輸入數據。右側是 cell 的內部結構。演算法旨在尋找能夠獲取準確網路的 cell。

在 AmoebaNet 中，谷歌所用進化演算法的一大重要特徵是採用了一種正則化形式：相比於移除最差的神經網路，他們移除了最老的神經網路（無論它有多好）。這提升了對任務優化時所發生變化的穩健性，並最終更可能得到更加準確的網路。

TPU 實現

項目地址：https://github.com/tensorflow/tpu/tree/master/models/experimental/amoeba_net

這些代碼是基於 AmoebaNet 論文的結果實現的，代碼的起始點來自 NASNet 實現和圖像處理代碼的分支。

預準備

1.建立一個谷歌雲項目

跟隨谷歌雲官網的 Quickstart Guide 中的指示來獲取 GCE VM（虛擬機）以訪問 Cloud TPU。

為了運行該模型，你需要：

• 一個 GCE VM 實例和相關的 Cloud TPU 資源；

• 一個 GCE bucket 來保存你的訓練檢查點；

• ImageNet 訓練和驗證數據預處理為 TFRecord 格式，並保存在 GCS（可選項）。

2.格式化數據

數據應該被格式化為 TFRecord 格式，可通過以下腳本完成：

https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py

如果你沒有準備好的 ImageNet 數據集，你可以用一個隨機生成的偽造數據集來測試模型。該數據集位於 Cloud Storage 的這個位置：gs://cloud-tpu-test-datasets/fake_imagenet。

訓練模型

通過執行以下命令來訓練模型（代入合適的值）：

python amoeba_net.py

--tpu_name=$TPU_NAME

--data_dir=$DATA_DIR

--model_dir=$MODEL_DIR

如果你沒有在相同的項目和 zone 中的 GCE VM 上運行這個腳本，你需要加上——project 和——zone 的 Flag 來為 Cloud TPU 指定你希望使用的值。

這將用單塊 Cloud TPU 在 ImageNet 數據集上以 256 的批量大小訓練一個 AMoebaNet-D 模型。通過默認的 Flag，模型應該能在 48 小時內（包括每經過幾個 epoch 之後的評估時間）達到超過 80% 的準確率。

你可以運行 TensorBoard（例如 tensorboard -logdir=$MODEL_DIR）來查看損失曲線和其它關於訓練過程的元數據。（注意：如果你在 VM 上運行，請確保恰當地配置 ssh 埠轉送或 GCE 防火牆規則。）

你也可以使用以下代碼在 7.5 小時內將 AmoebaNet-D 模型訓練到 93% 的 top-5 準確率。

python amoeba_net.py --tpu_name=$TPU_NAME

--data_dir=$DATA_DIR

--model_dir=$MODEL_DIR

--num_cells=6

--image_size=224

--num_epochs=35

--train_batch_size=1024

--eval_batch_size=1024

--lr=2.56

--lr_decay_value=0.88

--mode=train

--iterations_per_loop=1152

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！