Google本家から非公式(?)に出ているやつを使います。

https://github.com/tensorflow/haskell

基本的には、READMEに書いてあるようにstack dockerを使ってやれば良いのですが、このDockerfileではCPUでしか動かないので、GPUで動かしたい...という欲求がでてきました。ちなみに stack で docker を使うためには sudo 無しで動かす例の設定変更が必要です(私は変更後に再起動が必要でした)

https://github.com/commercialhaskell/stack/blob/master/doc/docker_integration.md

Dockerfileの変更

tensorflow-haskellのDockerfileでは、CPU用のイメージとビルド済みライブラリを使っています。書き換えた後のdiff はこんな感じで、

3c3
< FROM tensorflow/tensorflow:1.0.0
---
> FROM tensorflow/tensorflow:1.0.0-gpu
28,29c28,29
<     curl -O https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-cpu-linux-x86_64-1.0.0.tar.gz && \
<     tar zxf libtensorflow-cpu-linux-x86_64-1.0.0.tar.gz -C /usr/local && \
---
>     curl -O https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-gpu-linux-x86_64-1.0.0.tar.gz && \
>     tar zxf libtensorflow-gpu-linux-x86_64-1.0.0.tar.gz -C /usr/local && \

いろんなところに-gpuを足すだけです。ver1.0.0から変わっていたら合わせてください。あとはREADMEにあるようにDocker Imageをビルドします。

$ IMAGE_NAME=tensorflow/haskell:v0
$ docker build -t $IMAGE_NAME docker

無事Docker Imageができたので、このまま README にあるように MNIST を動かすと

$ cd ./tensorflow-mnist
$ stack --docker --docker-image=$IMAGE_NAME build --exec Main

以下のようなエラーがでて

E tensorflow/stream_executor/cuda/cuda_driver.cc:509] failed call to cuInit: CUDA_ERROR_NO_DEVICE
I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:152] no NVIDIA GPU device is present: /dev/nvidia0 does not exist

CPUで実行されてしまいました。これはGPU デバイスが見つからなかったということです。

nvidia-docker の準備

というわけで、nvidia-dockerを使います。nvidia-docker は docker 上のコンテナ(仮想OS)で、実OS側のGPUを使える凄いやつです。

https://github.com/NVIDIA/nvidia-docker

やはりREADMEに書いてあるような数行のインストールをします。nvidia-docker run --rm nvidia/cuda nvidia-smiでホスト側のデバイスがでてくれば成功です。

nvidia-docker の使い方

こちらの解説 「nvidia-dockerは何をしているのか」 によると, nvidia-docker は docker コマンドに対していい感じに引数を作って渡してるだけみたいです。というわけで、stack の `--docker-run-argsを使うと nvidia-docker 相当のことができます。

stack --docker --docker-run-args "$(curl -s http://localhost:3476/docker/cli)" --docker-image=$IMAGE_NAME build --exec Main

無事、実マシン側のGPUが認識されてGPUで実行しています。

I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:910] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
I tensorflow/core/common_runtime/gpu/gpu_device.cc:885] Found device 0 with properties: 
name: GeForce GTX 760
major: 3 minor: 0 memoryClockRate (GHz) 1.124
pciBusID 0000:01:00.0
Total memory: 1.95GiB
Free memory: 1.07GiB
I tensorflow/core/common_runtime/gpu/gpu_device.cc:906] DMA: 0 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:916] 0:   Y 
I tensorflow/core/common_runtime/gpu/gpu_device.cc:975] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 760, pci bus id: 0000:01:00.0)

GPUをオフにしたい

GPUを使わないときなど、一々dockerイメージを切り替えるのも不便です。 Tensorflow自体あまり詳しくないのですが、CUDA_VISIBLE_DEVICES環境変数をNO_DEVICEにセットすることで、GPUの利用を禁止できます。 stack の引数に --docker-env CUDA_VISIBLE_DEVICES=NO_DEVICE を設定すればOKです。

それにしてもstackとnvidia-dockerのおかげで、GHCやTensorflowだけでなくCUDAやCUDNNのバージョン管理もライブラリ開発者と同じものを簡単に再現できて便利になってきましたね。 それでは以上です。気が向けばAPIについてもメモしようと思います。