GPU 版 TensorFlow 環境を AWS 上にコマンド一発で立ち上げるスクリプトを書きました

こんにちは。吉田弁二郎です。

前置き

最近、TensorFlow のインストールは pip により簡単にできるようになりました。一方で GPU 環境の整備（CuDNN の準備等）が相変わらず面倒だと思うことも多く、GPU 自体の進歩も早いのでオンプレ環境に手を出すのも少しためらわれる感じもします。そこで Amazon Web Services (AWS) の GPU インスタンスを使ってみるものの、毎回ゼロから手作業で TensorFlow 環境を構築するのもそろそろ効率的にしたい、という状況になっています。

そんな訳で、コマンド一発で環境構築するための自動化スクリプトを書きました。もしかしたら誰かがやっているかもしれませんが気にしない。

やり方としては、Ansible を使います。Ansible はエージェントレスな構成管理ツールなどと言われ Web 系の人たちが使っていますが、要するにサーバ上でのディレクトリ作成やソフトウェアのインストール等を自動化するものです。リモートマシン、ローカルマシンのそれぞれに Python2 が入っていれば pip install ansible で使い始めることができます*1。

ここでは、Ubuntu Server 16.04 LTS (HVM) を選択し（＝あらかじめ GPU 開発環境が設定されていない）、最小限の工数で済ませるために nvidia-docker を使い、TensorFlow の docker image を取得することを大きな方針とします。

(2017.06.19 追記)
こんな AMI が登場したので、もはやこれでよさそうです。笑
AMI 自体が有料ではありますが（本記事の方法では、AMI は無料となります）。
AWS Marketplace: Deep Learning AMI Amazon Linux Version

本題

コード一式は
blog/misc/playbooks/tensorflow_on_aws_with_gpu at master · dlnp2/blog · GitHub
に置いたので、ご自由にお使いください。ここでは使い方を書いておきます。

今回は p2.xlarge インスタンス上に Ubuntu16.04 (xenial) を載せて作業します。インスタンスの起動、セキュリティグループの設定などをした後、まず SSH 接続できるか確認しましょう。

ssh -i ~/.ssh/private_key ubuntu@<instance ip here>

private_key は、適宜ご自分の環境で使っているものに変更してください。IP は何でもよいのですが、後の作業を楽にするためには Elastic IP を振っておいた方がいいです。続いてリモートマシンに SSH ログインした状態で Python 2.7 を入れます。

sudo apt-get update
sudo apt-get install python

「コマンド一発」じゃなかったのかよ！みたいに思った方、一旦、落ち着いてください。サーバでの作業は以上です。

Ansible を動かす

次に、ローカルマシンで Ansible の疎通確認を行います。まずコード一式を上記 GitHub レポジトリよりクローンしたら、hosts というファイルにインスタンス IP を記入してください。IP は複数でも大丈夫で、各々に対して同じ作業が自動実行されます。そして次を実行（ローカルマシンの Python は2系で、pip install ansible しておいてください）。

ansible -i hosts aws -m ping --user=ubuntu --private-key=~/.ssh/private_key

これで pong

your.ip.address.here | SUCCESS => {
    "changed": false,
    "ping": "pong"
}

が返ってくれば問題なし、エラーが出たら解消しましょう。最後に、

ansible-playbook -i hosts site.yml --user=ubuntu --private-key=~/.ssh/private_key

とすれば、CUDA のインストール、Docker 環境の構築、nvidia-docker の取得、そして TensorFlow image の取得が全て行われます。念のためリモートマシンに SSH ログインして

sudo nvidia-docker run -it gcr.io/tensorflow/tensorflow:latest-gpu /bin/bash

とした後に python コンソールで

import tensorflow as tf
hello = tf.constant('Hello TF!')
sess = tf.Session()
print sess.run(hello)

として Hello TF! と返ってくるようであれば、環境構築は無事成功です。後は、煮るなり焼くなり、好きに使っていきましょう。

なお、Ubuntu 14.04 で動かす場合は、/roles/docker/tasks/main.yml 内の "xenial" と書いてある部分を "trusty" に置き換える必要があると思います。（まだ確認してませんが、、）全面的にリファクタし、この部分は自動的に読みとるようにしました。（2017.04.09追記）

動作確認環境

• TensorFlow: 1.1.0-rc0
• Python 2.7.12（AWS）／Python 2.7.13（ローカルマシン）
• Ansible: 2.2.1.0
• AWS: p2.xlarge インスタンス（Tesla K80）
• AMI: ubuntu/images/hvm-ssd/ubuntu-xenial-16.04-amd64-server-20170221 (ami-a58d0dc5)
• Docker version 17.03.1-ce, build c6d412e（AWS）
• CUDA: 8.0

参考

• CUDA Toolkit Download | NVIDIA Developer
• Choose Docker CE or Docker EE on Ubuntu | Docker Documentation
• GitHub - NVIDIA/nvidia-docker: Build and run Docker containers leveraging NVIDIA GPUs
• Installing TensorFlow on Ubuntu  |  TensorFlow
• Ansible Documentation

*1:Python3 対応もされ始めているものの、現状は2系で安定稼働します。今回書いたコードも Python2 で動かすことを前提としています。ただ、Ansible を動かすために2系が必要というだけなので、数値計算用に3系が独立して存在していても問題ないはずです。