MXNet開放支持Keras，高效實現CNN與RNN的分散式訓練

知識 05-22

選自AWS Machine Learning Blog

作者：Lai Wei、Kalyanee Chendke、Aaron Markham、Sandeep Krishnamurthy

機器之心編譯

參與：路、王淑婷

今日 AWS 發布博客宣布 Apache MXNet 已經支持 Keras 2，開發者可以使用 Keras-MXNet 深度學習後端進行 CNN 和 RNN 的訓練，安裝簡便，速度提升，同時支持保存 MXNet 模型。

感謝 Keras 和 Apache MXNet（孵化）開源項目的貢獻者，Keras-MXNet 深度學習後端現已可用。

地址：https://github.com/awslabs/keras-apache-mxnet

Keras 是用 Python 編寫的高級神經網路 API，因其快速、簡單的 CNN 和 RNN 原型而廣受歡迎。

Keras 開發人員現在可以使用高性能的 MXNet 深度學習引擎展開卷積神經網路（CNN）和循環神經網路（RNN）的分散式訓練。通過更新幾行代碼，Keras 開發人員可以利用 MXNet 的多 GPU 分散式訓練功能來提高訓練速度。保存 MXNet 模型是該版本的另一個寶貴功能。你可以在 Keras 中進行設計，利用 Keras-MXNet 進行訓練，並使用 MXNet 在生產中運行大規模推斷。

用 Keras 2 和 MXNet 進行分散式訓練

本文介紹了如何安裝 Keras-MXNet，以及如何訓練 CNN 和 RNN。如果你以前使用過其它深度學習引擎進行分散式訓練，那你可能了解其中的難度和無趣。本文將展示如何使用 Keras-MXNet 進行訓練。

安裝僅需幾步

1. 部署 AWS 深度學習 AMI

2. 安裝 Keras-MXNet

3. 配置 Keras-MXNet

1. 部署 AWS 深度學習 AMI

按照 AWS 深度學習 AMI（DLAMI）部署教程來操作，教程地址：https://aws.amazon.com/getting-started/tutorials/get-started-dlami/。要利用多 GPU 訓練示例，請啟用 p3.8xlarge 或類似的多 GPU 實例類型。

想安裝用於運行 CUDA、Keras、MXNet 和其他框架（如 TensorFlow）的依賴項？請根據 Keras-MXNet 安裝指南進行操作：https://github.com/awslabs/keras-apache-mxnet/blob/master/docs/mxnet_backend/installation.md。

2. 安裝 Keras-MXNet

在 DLAMI 上的 MXnet Conda 環境中安裝 Keras-MXnet 及其依賴項。它已經包含 Keras 1.0 版，因此你需要先卸載此版本。登錄到 DLAMI 並運行以下命令：

Keras-MXnet 及其依賴項現已安裝在 DLAMI 上的 MXnet Conda 環境中。

3. 驗證 Keras-MXNet 安裝是否正確

使用以下代碼驗證 Keras 是否正在運行 MXNet 後端：

支持 CNN

現在我們在 CIFAR-10 數據集上訓練 ResNet 模型，來識別 10 個類別：飛機、汽車、鳥、貓、鹿、狗、青蛙、馬、船和卡車。我們可以使用 Keras-MXNet repo 示例部分中的 Keras 2 示例腳本（https://github.com/awslabs/keras-apache-mxnet/blob/master/examples/cifar10_resnet_multi_gpu.py）。使用 MXNet 作為 Keras 的後端幾乎不需要對腳本進行什麼更新。

首先，從 Keras-MXNet repo 文件夾中下載示例腳本：

腳本啟用 multi_gpu_model API，並輸入要使用的 GPU 數量。

然後，在終端窗口運行 nvidia-smi，以確定 DLAMI 上可用 GPU 的數量。接下來，如果你有四個 GPU，那麼你需要原樣運行該腳本；否則，你需要運行以下命令來打開腳本進行編輯：

該腳本具備以下行，該行定義 GPU 的數量。如有必要，更新該行。

現在，開始訓練。

（可選）訓練過程中，使用 nvidia-smi 命令檢查 GPU 利用和內存使用。打開另一個終端會話。

支持 RNN

Keras-MXNet 目前提供對 RNN 的實驗性支持。在使用 RNN 和 MXNet 後端時有一些局限性。更多信息，請查看 Keras-MXNet 文檔：https://github.com/awslabs/keras-apache-mxnet/blob/master/docs/mxnet_backend/using_rnn_with_mxnet_backend.md。此處的示例包括使用 LSTM 層訓練 IMDB 數據集時需要的一些變通方案。儘管有這些方案，但在多 GPU AMI 上訓練 RNN 會比你之前的訓練經驗容易一些，速度也更快。

使用 imdb_lstm 示例腳本：https://github.com/awslabs/keras-apache-mxnet/blob/master/examples/imdb_lstm.py。將輸入長度傳輸到嵌入層，按以下說明設置 unroll=True。

首先，在 DLAMI 的終端會話中，從 Keras-MXNet repo 文件夾中下載示例腳本：

然後，打開腳本，跳至下列行進行 review：

現在，示例腳本已經過修改，可與 MXNet 後端兼容，你可以運行以下行：

（可選）訓練過程中使用 nvidia-smi 命令檢查 GPU 利用和內存使用。打開另一個終端會話進行該操作。

基準

為了幫助大家評估不同 Keras 後端的性能，AWS 向 Keras-MXNet 添加了一個基準模塊。按表中描述在 CPU、單個 GPU 和多 GPU 上使用不同的模型和數據集，你會發現 Keras-MXNet 訓練 CNN 的速度更快，且在多個 GPU 上實現高效的性能提升。詳見訓練速度柱狀圖。關於如何運行基準腳本以及生成詳細的基準結果，詳見 Keras Benchmarks readme 文檔：https://github.com/awslabs/keras-apache-mxnet/tree/master/benchmark。

基準配置

Keras Version 2.1.6

MXNet Version 1.2.0

Image Data Format: Channel first

訓練 CIFAR10 數據集導致子線性縮放，因為該數據集圖像更小。CIFAR10 數據集包含 5 萬張圖像，每張圖像的大小是 32×32 像素。使用四個 GPU 傳輸這些小圖像所需的通信開銷比使用八個 GPU 的開銷高。ImageNet 和合成數據數據集更好地展示了 Keras-MXNet 可能帶來的性能改進。詳見下表。