PyTorch + JupyterLabのDocker環境構築

Deep Learningのフレームワークとして人気が上がっているPyTorchとデータ分析でPythonを使うときに人気のJupyterLab(Jupyter Notebookの後継)を使えるDocker環境の構築をしました。
新しく環境整備をしたので、記事を修正します(2019.12.14)

作業の流れ

NVIDIA GPUグラフックボードのドライバーやNVIDIA Container Toolkitをインストールする
JupyterLabのDockerfileに手を加え、Docker Imageを作成する
PyTorch公式のGitHubをもってくる
PyTorch公式のDockerfileに2.で作ったDocker Imageをベースに指定するなど必要な変更を加える。
PyTorchのDocker Imageをビルドする

具体的な手順

NVIDIA ドライバーやNVIDIA Container Toolkitをインストール

こちらの記事を参考にさせていただきました。かつて、Linuxマシンにグラフィックカードのドライバーのインストール、CUDA, cudnnのインストールを直接行ったことがありますが、Deep Learningのフレームワークとそれぞれのバージョンの組み合わせが違うとうまく動作しなくて苦戦しました。それに比べるとすごく楽になったと感じます。

NVIDIA Graphics Driverのインストール

ドライバのレポジトリをaptに登録します。

$ sudo add-apt-repository ppa:graphics-drivers/ppa
$ sudo apt update

推奨のドライバをインストールします。

$ sudo apt -y install ubuntu-drivers-common
$ sudo ubuntu-drivers autoinstall

NVIDIA Container Toolkitをインストールする

DockerでNVIDIAのGPUを使うのに必要なランタイムなどが含まれるNVIDIA Container Toolkitをインストールします。
はじめにリポジトリをaptに登録します。

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

そしてツールキットをインストールします。

$ sudo apt -y install nvidia-container-toolkit

一旦マシンを再起動しておきます。

$ sudo shutdown -r now

その後、下のコマンドでGPUが認識できているか確認できます。

$ nvidia-container-cli info

JupyterLabのベースとなるDockerfileを取得

JupterのGitHubをクローンしてきて、ベースとするDockerファイルを取得します。

$ git clone https://github.com/jupyter/docker-stacks.git

使用するファイル
- base-notebook/Dockerfile

JupyterLabのベースとするDockerfileへの変更を加える

base-notebook/DockerfileをビルドするときのベースをNVIDIAのDockerとするように変更します。
# の行はもともとの記述でコメントアウトで無効にして、その後の行を有効にしています。
私はbase-notebook/Dockerfileをテキストエディタで開き、冒頭の記述を以下のように変えました。
NVIDIAのDocker Hubのページを参考に、自分の使いたいDeepLearningフレームワークにあったバージョンを選んでください。

 #ARG BASE_CONTAINER=ubuntu:bionic-20191029@sha256:6e9f67fa63b0323e9a1e587fd71c561ba48a034504fb804fd26fd8800039835d
 #FROM $BASE_CONTAINER
 FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04

ベースDockerfileのビルド

base-notebookディレクトリで下のようなコマンドでDocker Imageを作成します。
-tの後ろは自由にDocker Imageの名前をつけて構いません。

$ docker image build ./ -t experiments/base-notebook

以下のコマンドでDocker Imageを表示し、作成できたか確かめます。

$ docker images

Jupyter LabのDocker ImageをベースにPyTorchのDocker Imageを作成する。

自分の保存したいディレクトリで下のようなコマンドでPyTorch公式のGitHubをクローンします。

$ git clone https://github.com/pytorch/pytorch.git

PyTorchのDockerfileに変更を加える

docker/pytorch/Dockerfileをdocker/pytorch-notebook/Dockerfileとしてコピーし、必要な変更を行います。
/pytorch-notebook/Dockerfileをテキストエディタで開き、冒頭を以下のように変更して前の手順で作成したベースとするJupyter LabのDocker Imageをベースにするようにします。

#FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu16.04
FROM experiments/base-notebook:latest

PyTorchをインストールする前にminiconda(Anacondaの軽量版)をインストールする箇所がありますが、
Jupyer LabのDockerでインストールしているため、コメントアウトで無効化し、その他のライブラリやpytorchをインストールするところから実行するようにしています。有効にする行の頭にはRUNをつけてください。
PyTorchのチュートリアルプログラムなどを実行して、以下のようなパッケージをインストールするように追記しています。
* ipykernel pandas matplotlib scikit-learn pillow seaborn tqdm ipywidgets
* opencv

* libxkbcommon-x11-0 (pyside2で作ったプログラムを実行時に必要)

 # Install PyTorch
 #RUN curl -o ~/miniconda.sh -O  https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
 #     chmod +x ~/miniconda.sh && \
 #     ~/miniconda.sh -b -p /opt/conda && \
 #     rm ~/miniconda.sh && \
 RUN  /opt/conda/bin/conda install -y python=$PYTHON_VERSION numpy pyyaml scipy ipython mkl mkl-include ninja cython typing \
    ipykernel pandas matplotlib scikit-learn pillow seaborn tqdm openpyxl ipywidgets && \
    /opt/conda/bin/conda install -y -c pytorch magma-cuda100 && \
    /opt/conda/bin/conda install -y -c conda-forge opencv pyside2 && \
    /opt/conda/bin/conda clean -ya
ENV PATH /opt/conda/bin:$PATH

追記：
opencvのインポート時に次のようなエラーが出ました。「ImportError: libGL.so.1: cannot open shared object fil e: No such file or directory」
apt-get installをするところで「libgl1-mesa-dev」を加えました。(この記事を参考)

JupyterLabのDockerのユーザー環境と合わせるため、終わりの方の下の記述をコメントアウトしました。
WORKDIR /workspace
RUN chmod -R a+w .
代わって、下の記述を加えました。

    RUN chown -R $NB_UID:$NB_GID /home/$NB_USER
    WORKDIR /home/$NB_USER
    # Switch back to jovyan to avoid accidental container runs as root
    USER $NB_UID
    RUN echo 'export PATH=/opt/conda/bin:$PATH'>> ~/.bashrc

PyTorchのDocker Imageを作成する

GitHubからcloneしてきたPyTorchの"ルートディレクトリ"(ここは結構間違えやすいので注意してください。GitHubからサブモジュールをアップデートし、cmakeなどするため決められた位置である必要がある）で、下のようなコマンドでDocker Imageをビルドします。
今回の例では、作成するDocker Imageの名前を”experiments/pytorch-notebook"として出力します。


$ docker build -t experiments/pytorch-notebook -f docker/pytorch-notebook/Dockerfile .

注意点として、caffe2のためのcmakeの処理がかなり時間がかかります。

作成したDocker Imageを使用する

作成したDocker Imageからコンテナを作って実行します。
最初にJupyter Labにブラウザでアクセスするときのパスワードを設定します。
こちらの記事を参考にさせていただきました。

docker run \
 --rm -it \
 --user root \
 --name pytorch-notebook \
 experiments/pytorch-notebook:latest \
 /bin/bash -c \
 "python -c 'from notebook.auth import passwd;print(passwd())'"

パスワードを入力するよう促されるので、２度入力します。
ハッシュ化されたパスワードの値(sha1:xxxxxxxxxxxxxxxxxxxxxxxx)が出力されるので、記録しておきます。

Enter password:
Verify password:
sha1:xxxxxxxxxxxxxxxxxxxxxxxx

ハッシュ化されたパスワードを付与して(--NotebookApp.password=に指定する)Jupyter Labを起動します。

docker run \
 --rm \
 --user root -e NB_UID=$UID \
 -p 58888:8888 -p 50022:22 -p 56006:6006 \
 -v ~/:/home/jovyan/work \
 --name pytorch-notebook \
 --gpus all \
 --ipc=host \
 experiments/pytorch-notebook:latest \
 start.sh jupyter lab --NotebookApp.password="sha1:xxxxxxxxxxxxxxxxxxxxxxxx"

Webブラウザでlocalhost:58888(上の例のようにポート番号をマッピングした場合)にアクセスするとJupyter Labが使えます。

PyTorchでGPUを使うときは、--ipc=host か --shm-size=16G のようなオプションをつけてメモリ割当をしておく必要があるようです。
ミニバッチ作成時のDataLoaderでnum_workersを1以上に設定して、マルチプロセスを利用すると、共有メモリを用いてデータがやり取りされることに起因すると思われます。Qiitaの参考記事

pythonファイルを実行する場合は、%runを使います。

%run -i sample.py

References

[1] PyTorch GitHub
[2] Jupyte Lab Dockerfile
[3] NVIDIA Container Toolkit を使って Docker コンテナで GPU を使う
[4] Jupyter LabをDockerで環境構築する

PyTorchとJupyterLabが使えるDocker環境を構築する