作者：Azalia Mirhoseini等機器之心編譯參與：吳攀、李澤南
眾所周知，深度學習是非常計算密集的，合理分配計算資源對於提升運算速度、節省能源等有著至關重要的價值。近日，谷歌的一篇 ICML 2017 論文提出了一種使用強化學習優化設備配置的方法，據稱可以在一些網絡計算上實現優於專家優化的方法的表現。機器之心對本文進行了摘要介紹。
論文：使用強化學習的設備配置優化（Device Placement Optimization with Reinforcement Learning）
論文地址：https://arxiv.org/abs/1706.04972
摘要
過去幾年，我們見證了神經網絡的訓練和推理推動著計算能力需求的增長。目前解決這種需求的常用辦法是搭建異構分布式計算環境，其中包含 CPU 和 GPU 這樣的硬體設備。重要的是，神經網絡模型在設備上的計算資源分配通常是由人類專家依靠簡單的探索和直覺做出的。在本文中，我們提出了一種為 TensorFlow 計算圖學習並優化設備任務分配的方法。其中最關鍵的思路是使用序列到序列模型來預測 TensorFlow 圖中哪些子集應該運行在哪些可用設備上，然後將預測分配的執行時間用作獎勵信號以優化序列到序列模型的參數。我們的結果在用於 ImageNet 分類的 Inception-V3 任務、RNN LSTM 語言建模和神經機器翻譯任務上的表現優於普通手動設備配置（device placement）方式。
圖 1 ：基於強化學習的設備配置模型概覽
圖 2：設備配置模型的架構
3.2 架構細節
我們使用了一種序列到序列模型（Sutskever et al., 2014）來預測設備配置，該模型帶有 LSTM（Hochreiter & Schmidhuber, 1997）和一個基於內容的注意機制（Bahdanau et al., 2015）。圖 2 展示了我們模型的整體架構，其可以分成兩個部分：編碼器 RNN 和解碼器 RNN。
編碼器 RNN 的輸入是輸入圖（input graph）的指令的序列。我們通過連接它們的信息嵌入這些指令。具體來說，對於每一個圖 G，我們首先收集其指令的類型。其指令的類型描述了其基本的計算，比如 MatMul 或 conv2d。對於每一種類型，我們都存儲一個可調節的嵌入向量。然後我們記錄輸出張量的每個指令列表的大小，然後把它們連接成一個固定大小的填充了 0 的列表，稱為輸出形狀（output shape）。我們也取其 one-hot 編碼向量，其表徵了那些將輸入和輸出指引到每個指令的指令。最後，每個指令的嵌入都是其類型、其輸出形狀和其 one-hot 編碼的鄰接信息的連接。
該解碼器是一個帶有固定時間步驟數量的 attentional LSTM（Bahdanau et al., 2015），這個數量等於圖 G 中的指令的數量。在每一步，該解碼器輸出該設備（device）以運行在同一編碼器時間步驟的指令。每個設備都有其自己的可調節嵌入，然後其被用作下一個解碼器時間步驟的輸入。
圖 3：分布式和異步的參數更新與獎勵評估
表 2：由基於強化學習的方法所找到的設備配置方式和基準方法的運行時間比較（單位：秒，越低越好）。
對於每一個模型，第一行給出了帶有 1 個 CPU 和 2 個 GPU 的結果；第二行給出了帶有 1 個 CPU 和 4 個 GPU 的結果。最後一列給出了基於強化學習的配置優於最快基準方法的運行時間的提升。為了減少方差，少於 10 秒的運行時間都測試了 15 次，這裡給出的是平均值。OOM 是指內存不足（Out Of Memory）。
圖 4：神經機器翻譯圖的基於強化學習的配置。上面是編碼器，下面是解碼器。設備用不同顏色表示，其中無色表示指令在 CPU 上運行，不同的其它顏色表示不同的 GPU。相比於專家精心設計的配置，這個配置在運行時間上實現了 19.3% 的提升。
圖 5：Inception-V3 的基於強化學習的配置。設備用不同顏色表示，其中無色表示指令在 CPU 上運行，不同的其它顏色表示不同的 GPU。相比於專家設計的配置，這個配置在運行時間上實現了 19.7% 的提升