Edited at

Dockerで機械学習する環境構築


動機&概要

機械学習の環境構築は面倒ですよね。盛大に環境も汚染しますし。というわけで最近ではイケてる機械学習エンジニア(n = 3くらい)はDockerに機械学習用のあれこれを閉じ込めています。

機械学習をする時にはJupyterを使う人が多いと思います。ところでここに Jupyter Docker Stacksという素晴らしいプロジェクトがあります。

https://github.com/jupyter/docker-stacks

こちらはJupyterを起動可能なDockerがまとまっている素晴らしいレポジトリです。しかしこちらのDocker Stacksには多少の不便があります。


  • GPU対応していない。TensorFlowも入っているけれど、CPU版で回すしかない

  • Dockerファイルだけなので、起動コマンドなどがめんどくさい


作ったやつ

というわけで、上記2点を解決したものを作りました。

https://github.com/tackman/docker-stacks/tree/gpu-support


前提


  • NVIDIAのグラフィックカード搭載をしているLinux環境を想定しています

  • ドライバ・Docker・Nvidia Docker(2系)はインストール済み

CUDAや各種Pythonライブラリのインストールは不要です(それらはDocker内で処理されます)


できること&できないこと


  • Jupyter上からPytorchでGPUを利用して学習スクリプトなどを回せます


  • TensorFlowも一応インストールされていますがCPU版しか動きません・・・・


    • 個人的にあまりTensorFlowを直接触りたくないのでやる気湧かなかった。TensorFlowでCUDA使いたい人は頑張ってください



  • 使うだけなら docker-compose.yml をJupyter起動したいディレクトリにコピーして、docker-compose upするだけでJupyterLabが走ります

  • 入っているライブラリはTensorFlow / Keras, Pytorch (+ もともとJupyter Docker Stacksで入っているもの。scipyとか必要そうなものは一通り入っています)


中身の簡単な説明

Jupyter Docker Stacksで提供されているDockerファイルを、大本になっているDockerイメージをNVIDIA提供のCUDA / cuDNN入ってる版に差し替え。合わせて起動をNvidia Dockerに合わせただけ。あと個人的に使いたかったPytorchインストールを追加ました。


まとめ

今回紹介したDockerイメージを使うかは別として、機械学習向けの環境汚染で死んだ経験のある人はDocker採用おすすめです。ドライバインストールという最大の闇は解決しませんが、それ以外は文明化して生きていけます。


追記 2018.10.12

TensorFlowのインストール方法を変更したところ、GPUで動くようになりました。