Bash on Windows上でTensorFlowをBuildする(CPU向け SSE/AVX対応)

Bash on Windows上で、SSE/AVXの拡張命令付きでTensorFlowをBuildしたので、その時のメモ。
結論としては(時間はかかるが)非常に簡単だった。

環境

• Surface Pro (2017)
  • Core i7-7660U(2.50/4.00GHz)
  • 16GB RAM
  • 512GB SSD
  • Win10 Pro Creators Update
• miniconda v4.3.22
  • Linux x64用を公式手順でインストール。
• Python=3.6.1
• numpy=1.13.0
  • 事前に conda install でインストールしておく。

pip install したTensorFlowはSSE/AVXといった拡張命令に対応していないため、以下のようなWarningが表示される。

2017-06-25 14:48:26.653110: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\tensorflow\core\platform\cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE instructions, but these are available on your machine and could speed up CPU computations.
2017-06-25 14:48:26.653282: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\tensorflow\core\platform\cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE2 instructions, but these are available on your machine and could speed up CPU computations.
2017-06-25 14:48:26.653435: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\tensorflow\core\platform\cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE3 instructions, but these are available on your machine and could speed up CPU computations.
2017-06-25 14:48:26.654259: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\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.
2017-06-25 14:48:26.654920: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\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.
2017-06-25 14:48:26.655425: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\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.
2017-06-25 14:48:26.655783: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\tensorflow\core\platform\cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX2 instructions, but these are available on your machine and could speed up CPU computations.
2017-06-25 14:48:26.656121: W c:\tf_jenkins\home\workspace\release-win\m\windows\py\36\tensorflow\core\platform\cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use FMA instructions, but these are available on your machine and could speed up CPU computations.

せっかくなので、拡張命令に対応したビルドを作ってみることにした。
以下、minicondaの仮想環境上で作業する。

Bazelのセットアップ

Bash on Windows上でも、Ubuntu向けの公式手順 に従えばよい。
以下、2017/06時点でのインストール手順。

$ echo "deb [arch=amd64] http://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
$ curl https://bazel.build/bazel-release.pub.gpg | sudo apt-key add -
$ sudo apt-get update && sudo apt-get install bazel

あとは待っているだけでBazelのセットアップが完了する。非常に簡単。

TensorFlowのビルド

事前準備はここまでで、早速本体のビルドに入る。
まずは、公式リポジトリからソースコードを git clone し、任意のバージョンをCheckoutする。今回は v1.2.0 のタグで作業を行った。

Configure

./configure を実行する。今回は下記のオプションを利用した。

$ ./configure
You have bazel 0.5.1 installed.
Please specify the location of python. [Default is /home/xxx/miniconda3/envs/xxx/bin/python]:
Found possible Python library paths:
  /home/xxx/miniconda3/envs/xxx/lib/python3.6/site-packages
Please input the desired Python library path to use.  Default is [/home/xxx/miniconda3/envs/xxx/lib/python3.6/site-packages]
Using python library path: /home/xxx/miniconda3/envs/xxx/lib/python3.6/site-packages
Do you wish to build TensorFlow with MKL support? [y/N] y
MKL support will be enabled for TensorFlow
Do you wish to download MKL LIB from the web? [Y/n] Y
Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native]:
Do you wish to use jemalloc as the malloc implementation? [Y/n] Y
jemalloc enabled
Do you wish to build TensorFlow with Google Cloud Platform support? [y/N] N
No Google Cloud Platform support will be enabled for TensorFlow
Do you wish to build TensorFlow with Hadoop File System support? [y/N] 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] N
No XLA JIT support will be enabled for TensorFlow
Do you wish to build TensorFlow with VERBS support? [y/N] N
No VERBS support will be enabled for TensorFlow
Do you wish to build TensorFlow with OpenCL support? [y/N] N
No OpenCL support will be enabled for TensorFlow
Do you wish to build TensorFlow with CUDA support? [y/N] N
No CUDA support will be enabled for TensorFlow
Do you wish to build TensorFlow with MPI support? [y/N] N
MPI support will not be enabled for TensorFlow
Configuration finished

ビルド実行

以下のコマンドでBuldを実行する。
--copt=xxx の部分で有効化する拡張命令を設定する。対象のCPUに合わせて適宜変更が必要。
今回は下記を有効化した。
* AVX2
* AVX
* SSE4.1
* SSE4.2
* FMA

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

このまま1時間程度待てば、ビルドが完了する。
尚、このときSurface Pro本体が非常に高温になった
時間がかかることも併せて、膝の上などで実行するのはお勧めできない。

whlの生成

ビルドが終わったら、最後にwhlファイルを生成する。
下記のコマンドのうち、/tmp/tensorflow_pkg 部分を変更すれば、任意の場所にファイルを出力可能。

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

ここまで終われば、あとは生成されたwhlファイルを pip install すればよい。

参考までに、この環境では実行速度が20～30%程度高速化した。