nvidia-dockerでtensorflowを動かす

はじめに

前回の 記事 ではtensorflow/tensorflowのイメージを元にdocker上でGPUモードのtensorflowを動かせる環境を構築しました．
先日，caffeのdocker環境を構築する機会がありまして，nvidia-dockerというものを使いました．nvidia-dockerを使うとdocker上から簡単にcudaにアクセスすることができます．今回はnvidia-dockerを使ってtensorflowを動かしてみようと思います．

追記 - 現在はtensorflow公式のイメージもnvidia-dockerを使うようになっています．

nvidia-dockerの導入

dockerを導入した環境において，以下で導入できます．

$ git clone https://github.com/NVIDIA/nvidia-docker.git
$ cd nvidia-docker
$ make install
$ nvidia-docker volume setup

Dockerfile

ベースイメージとしてnvidia/cuda:cudnnを指定しています．バージョンを指定したほうが安心かもしれませんが．cuda，cudnnが導入されたイメージです．
pip install --upgrade https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-0.8.0-cp27-none-linux_x86_64.whlは普通の環境と同様の方法で導入しています．
tensorflow/のイメージだとjupyterが入っていたりしますが，個人的に使っていないので．

FROM nvidia/cuda:cudnn

RUN apt-get update
RUN apt-get install -y --no-install-recommends python-pip python-dev

RUN pip install --upgrade https://storage.googleapis.com/tensorflow/linux/gpu/tensorflow-0.8.0-cp27-none-linux_x86_64.whl

ENTRYPOINT ["/bin/bash"]

起動

dockerの代わりにnvidia-dockerで起動します．device等を指定しなくてもいいので楽です．

nvidia-docker run -it image_name