Python: Keras/TensorFlow の学習を CPU の拡張命令で高速化する (Mac OS X)

今回のネタは TensorFlow を使っていると、いつも目にしていた警告について。

それは、次のようなもの。

W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.1 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.2 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX instructions, but these are available on your machine and could speed up CPU computations.

どうやら TensorFlow は CPU の拡張命令を使うことができるらしい。しかし、PyPI 経由で一般に配布されているバイナリパッケージでは、それらは有効になっていない。もちろん、同じ x86 系 CPU であってもどんな拡張命令セットを持っているかは製品によって異なるので、これは当然の判断だろう。

そこで、今回は CPU の拡張命令を有効にして TensorFlow をソースコードからコンパイルしてみることにした。そうすることで、どれくらい学習を高速化できるだろうか？というのが主題となる。

先に結論から書いてしまうと、今回使った環境では大体 30% くらい学習が速くなった。ぶっちゃけ、これくらいであれば素直に GPU を使える環境を用意した方が良いと思う。

blog.amedama.jp

ただし、自分で使う環境向けにチューンするという面において、自前でコンパイルするのは良い選択肢かもしれないと感じた。

今回使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.12.3
BuildVersion:   16D32
$ python --version
Python 3.6.0

あらかじめ、計算速度のベンチマークに使うアプリケーションをダウンロードしておこう。

$ curl -O https://raw.githubusercontent.com/fchollet/keras/master/examples/mnist_cnn.py
$ echo 'K.clear_session()' >> mnist_cnn.py

拡張命令を使わない場合

まずは、CPU の拡張命令を使わない場合にかかる時間を調べておこう。

PyPI から汎用のバイナリパッケージで TensorFlow をインストールする。ベンチマークにするアプリケーションを実行するために Keras も入れておく。

$ pip install keras tensorflow

インストールできたら Python の REPL を起動する。

$ python

REPL が起動したら MNIST のデータセットをダウンロードしておこう。これには少し時間がかかる。

>>> from keras.datasets import mnist
Using TensorFlow backend.
>>> mnist.load_data()
Downloading data from https://s3.amazonaws.com/img-datasets/mnist.pkl.gz

終わったら REPL から抜けておく。

>>> exit()

それでは、まずは CPU の拡張命令を使わない場合にかかる時間を測ろう。 time コマンド経由でベンチマーク用のアプリケーションを実行する。

$ time python mnist_cnn.py
Using TensorFlow backend.
X_train shape: (60000, 28, 28, 1)
60000 train samples
10000 test samples
Train on 60000 samples, validate on 10000 samples
Epoch 1/12
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.1 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.2 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX instructions, but these are available on your machine and could speed up CPU computations.
60000/60000 [==============================] - 92s - loss: 0.3741 - acc: 0.8859 - val_loss: 0.0880 - val_acc: 0.9718
...(省略)...
python mnist_cnn.py  2452.87s user 859.80s system 302% cpu 18:15.11 total

CPU の拡張命令を使わない場合には 2450 秒かかった。

拡張命令を使う場合

次は CPU の拡張命令を使う場合について。警告メッセージを見る限り、今回使うマシンでは SSE4.1 SSE4.2 AVX という拡張命令が使える。

それらの拡張命令を有効にするためにソースコードから TensorFlow をインストールする。ソースコードからのインストールには次のドキュメントを参照すると良い。 Download and Setup | TensorFlow

まずは、Homebrew を使って依存パッケージをインストールする。

$ brew install bazel swig

次に TensorFlow のソースコードをクローンする。

$ git clone https://github.com/tensorflow/tensorflow.git
$ cd tensorflow

先ほどベンチマークしたのと同じバージョンのタグをチェックアウトしよう。

$ git checkout v1.0.1

次にインストールするための環境を configure で設定していく。今回は CUDA を使わないので、さほど難しくはない。

$ ./configure
Please specify the location of python. [Default is /Users/amedama/.virtualenvs/py36/bin/python]:
Please specify optimization flags to use during compilation [Default is -march=native]:
Do you wish to use jemalloc as the malloc implementation? (Linux only) [Y/n] n
jemalloc disabled on Linux
Do you wish to build TensorFlow with Google Cloud Platform support? [y/N]
No Google Cloud Platform support will be enabled for TensorFlow
Do you wish to build TensorFlow with Hadoop File System support? [y/N]
No Hadoop File System support will be enabled for TensorFlow
Do you wish to build TensorFlow with the XLA just-in-time compiler (experimental)? [y/N]
No XLA support will be enabled for TensorFlow
Found possible Python library paths:
  /Users/amedama/.virtualenvs/py36/lib/python3.6/site-packages
Please input the desired Python library path to use.  Default is [/Users/amedama/.virtualenvs/py36/lib/python3.6/site-packages]

Using python library path: /Users/amedama/.virtualenvs/py36/lib/python3.6/site-packages
Do you wish to build TensorFlow with OpenCL support? [y/N]
No OpenCL support will be enabled for TensorFlow
Do you wish to build TensorFlow with CUDA support? [y/N]
No CUDA support will be enabled for TensorFlow
Configuration finished
Extracting Bazel installation...
..........................................................
INFO: Starting clean (this may take a while). Consider using --expunge_async if the clean takes more than several minutes.
.........................................................
INFO: All external dependencies fetched successfully.

設定が終わったら次に bazel を使ってビルドする。ここで使用する CPU の拡張命令を指定する。 --copt=-m の後ろに拡張命令の名前を入力しよう。この処理には時間がかかるので気長に待つ。

$ bazel build -c opt --copt=-mavx --copt=-msse4.1 --copt=-msse4.2 //tensorflow/tools/pip_package:build_pip_package

もし、ここで指定しなければ CPU 拡張命令を使わない汎用なパッケージになる。

続けて Wheel パッケージをビルドする。二番目に渡している /tmp/tensorflow_pkg が Wheel パッケージを保存する場所になる。

$ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

これで CPU の拡張命令が有効になった Wheel パッケージができた。

$ ls /tmp/tensorflow_pkg
tensorflow-1.0.1-cp36-cp36m-macosx_10_12_x86_64.whl

既にインストールされている汎用なパッケージをアンインストールした上で、インストールしよう。

$ pip uninstall -y tensorflow
$ pip install /tmp/tensorflow_pkg/tensorflow-1.0.1-cp36-cp36m-macosx_10_12_x86_64.whl
$ pip list --format=columns | grep -i tensorflow
tensorflow       1.0.1

さて、先ほどと同じようにベンチマークとなるアプリケーションを実行してみよう。どれくらい速くなるだろうか。

$ time python mnist_cnn.py
Using TensorFlow backend.
X_train shape: (60000, 28, 28, 1)
60000 train samples
10000 test samples
Train on 60000 samples, validate on 10000 samples
Epoch 1/12
60000/60000 [==============================] - 75s - loss: 0.3735 - acc: 0.8859 - val_loss: 0.0888 - val_acc: 0.9724
...(省略)...
python mnist_cnn.py  1871.67s user 901.15s system 296% cpu 15:34.90 total

今度は 1870 秒で終わった。

結果は CPU の拡張命令を使わない場合が 2450 秒で、使った場合が 1870 秒となった。つまり、処理速度が 30% ほど速くなった。また、一回のエポックでいうと、だいたい 20% ほど速くなっているようだ。

まとめ

今回は、普段 TensorFlow を使っているときに表示される警告が気になってソースコードからビルドを試してみた。 PyPI で配布されている TensorFlow は CPU に依存しない汎用なバイナリになっている。もし、自分の CPU に合わせて拡張命令を使ったバイナリを作りたいときは、ソースコードからコンパイルする必要がある。 CPU の拡張命令を有効にすると、今回のケースでは 30% ほど学習が速くなった。 GPU を使った学習の高速化に比べると、これは微々たるもののように感じる。しかし、自分の環境に向けてチューンされたバイナリを作るという意味では、ソースコードからのコンパイルは有用かもしれない。

いじょう。