Dockerで機械学習環境を作ろう!

はじめに

こんにちは、技術開発の三浦です。最近開発用のノートPCを手に入れました。今回はそこにDockerで機械学習の環境を作った話を紹介します！

構成

機械学習関連のタスクはPythonで行っているのですが、その場の勢いでライブラリをインストールしたりしてしまい、いざアプリケーションに落とし込もうとすると「どうやって環境作ってたっけ・・・」と環境の再現に無駄に時間がかかっていました。あと、ホスト側に直で色々入れてしまうと何か問題があった時に戻すのが大変です(一応Pythonの仮想環境pyenvは使っているのですが)。そこでホスト側にはなるべく手を入れず、Dockerコンテナ上にPythonの実行環境を構築、Jupyter Notebookサーバを立ち上げてブラウザから操作できる環境を作りました。Dockerのイメージ構築手順はDockerfileに記載し、環境再現性も持たせるようにしました。

ディレクトリ構成は以下のようにしました。

ml_env/
　┣Dockerfile
　┣docker-compose.yml
　┣.env
　┣requirement.txt
　┣source/

ちなみにホストOSはUbuntu 18.04です。

Dockerの導入

開発用のPCにはNVIDIAのGPUが搭載されているので、DockerコンテナからGPUにアクセスできるようNVIDIA Container Toolkitを導入します。こちらの手順を参考にしてNVIDIA driver/Docker/NVIDIA Container Tooklitを導入しました。

github.com

Dockerfileの内容

Dockerfileは以下のようにしました。

FROM python:3.8
ARG USER
ENV HOME /home/${USER}
RUN apt update
RUN apt upgrade -y
RUN useradd -m ${USER}
RUN gpasswd -a ${USER} sudo
RUN echo "${USER}:pass" | chpasswd
ADD requirement.txt ${HOME}
WORKDIR ${HOME}
RUN pip install -r requirement.txt
USER ${USER}
WORKDIR ${HOME}/source

コンテナ内でrootユーザで動いているとホストから消せないファイルなど出来てしまうことがあって嫌なので別のユーザを追加しています。ユーザ名はbuild時の引数で指定するようにしています。Pythonのライブラリはrequirement.txtに指定するようにしています。例えばこんな感じです。

jupyter
jupyterthemes
scikit-learn
pandas
tensorflow-gpu
xgboost
opencv-python

Jupyter Notebookのデザインを変更できるようにjupyterthemesを導入しています。適用するテーマはコンテナ起動時に変更できるよう環境ファイル.envで指定するようにしています。

nvidia-container-runtimeの導入

docker-composeで立ち上げたコンテナからGPUにアクセスできるようnvidia-container-runtimeを導入します。こちらの手順に従いました。

github.com

curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | \
  sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-container-runtime/$distribution/nvidia-container-runtime.list | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-runtime.list
sudo apt-get update
sudo apt-get install nvidia-container-runtime

またdocker-composeでruntimeオプションがないといったエラーが出たのですが、docker-composeのバージョンが古かったようです。

$ docker-compose -v
docker-compose version 1.17.1, build unknown

上記の1.17.1ではエラーが出たため、こちらの手順に従って1.25.4にバージョンを上げました。

github.com

上記手順のあと、OSを再起動すると新しいバージョンを認識するようになりました。

docker-compose.yml

docker-composeは以下のようにしました。

version: '2.4'
services:
    jupyter:
        build:
            context: ./
            args: 
                - USER=${USER}
        volumes:
                - ./source:/home/${USER}/source
        ports:
                - '9999:9999'
        runtime: nvidia
        environment:
                - NVIDIA_VISIBLE_DEVICES=all
        command: 
                sh -c "jt -t ${THEME}; jupyter notebook --port 9999 --ip 0.0.0.0"

環境変数USERとTHEMEは.envファイルで指定します。たとえばこんな感じです。

USER=miu
THEME=monokai

(monokai、好きなのです。)

あとは以下のコマンドで環境のbuildとコンテナを立ち上げます。

$ docker-compose build
$ docker-compose up -d

ホストのブラウザでlocalhost:9999にアクセスするとコンテナ内で立ち上がってるJupyter Notebookにアクセスできます。またJupyter Notebookでターミナルを立ち上げてnvidia-smiコマンドを実行してGPUにアクセスできることも確認できます。

nvidia-smi実行結果

最後に

自分好みの環境を作るのは楽しいです(もちろん色々大変ですが・・・)。今回作った機械学習環境でこれからいろんなプロダクトを作っていきたいです！

2020/04/01追記

上記のdocker-compose.ymlで立ち上げたコンテナからcudaが使えなかったため、docker-compose.ymlのenvironmentを以下のように変更しました。

environment:
    - NVIDIA_VISIBLE_DEVICES=all
    - NVIDIA_DRIVER_CAPABILITIES=compute,utility

"NVIDIA_DRIVER_CAPABILITIES"を追加し、cudaを使用するためcomputeを、nvidia-smiを使用するためutilityを指定しています。