Pytorch+JupyterLab+GPUをDockerで作成

ぐぐりながら見様見真似でdockerを使ってみました。

目的

• PyTorchによる発展ディープラーニングができる環境を作る
• やるならJupyterLabを使いたい
• Dockerでのコンテナ化もやってみたい

結果

初めはjupyterのdocker-stacksのベースを書き換えたりして試行錯誤してましたが、公式のpytorch/pytorchをベースにjupyterlabなど必要なライブラリをインストールすることで目的のものができました。

環境

Dockerやnvidia-container-toolkitはインストール済みとします。
- GPU: GTX1080
- OS: Ubuntu 18.04.3 LTS
- Nvidia driver: 440.48.02 CUDA Version: 10.2
- Docker: 19.03.5

手順

適当な作業用フォルダを作成し移動

$ mkdir pytorch-lab[なんでも良い]
$ cd pytorch-lab

Dockerfileを作成

$ vim Dockerfile

Dockerfile

FROM pytorch/pytorch:latest

# Install required libraries
RUN conda config --add channels pytorch \
 && conda config --append channels conda-forge \
 && conda update --all --yes --quiet \
 && conda install --yes --quiet \
    ipywidgets \
    jupyterlab \
    matplotlib \
    nodejs \
    opencv \
    pandas \
    scikit-learn \
    seaborn \
    sympy \
 && conda clean --all -f -y

# Install jupyter extensions
RUN jupyter nbextension enable --py --sys-prefix widgetsnbextension \
 && jupyter labextension install @jupyter-widgets/jupyterlab-manager

COPY jupyter_notebook_config.py /root/.jupyter/

やっていることは

• base imageとしてpytorch/pytorchを使用。
• そのままconda updateを行うとpytorchがダウングレードされてしまうので、conda config --add channels pytorchでpytorchのライブラリを高優先度で追加。
• condaのライブラリだけだと足りないことがあるのでconda config --append channels conda-forgeでconda-forgeのライブラリを低優先度で登録。
• 必要なライブラリのインストール。
• jupyterlabでtqdmでの進捗バー表示が行えるようにjupyterlabの拡張機能をインストール。
• 後述するjupyterの設定ファイルをコピー。

jupyterの設定ファイルを作成。

Dockerfileと同じ場所にjupyter_notebook_config.pyを作成。

私は試行錯誤していた際に使用したjupyter/docker-stacks/base-notebook/jupyter_notebook_config.pyをコピーして設定を書き加えています。
追加した設定は

• c.NotebookApp.allow_root = True : rootでの起動を許可
• c.NotebookApp.password = 'パスワードのハッシュ値' : ログインパスワードを設定

パスワードのハッシュ値はローカルのpythonにipythonがインストールされていれば、以下のコマンドで生成できます。

$ ipython

In [1]: from IPython.lib import passwd
In [2]: passwd()

パスワードを設定して表示される'sha1:~'をコピーしておいてください。

jupyter_notebook_config.py

c.NotebookApp.allow_root = True
c.NotebookApp.password = 'sha1:~'

これらをc.NotebookApp.open_browser = Falseの後ろあたりに打ち込んでください。

imageのビルド

上記のDockerfileおよびJupiter_notebook_configのあるフォルダでビルドを行います。
pytorch-labのところはわかりやすい任意の名前をつけてください。
初回はベースのpytorch/pytorchのダウンロードが行われるのでしばらく時間がかかります。

$ docker build -t pytorch-lab ./

docker imagesで作成したイメージが追加されていたら成功です。

$ docker images
REPOSITORY              TAG                      IMAGE ID            CREATED             SIZE
pytorch-lab             latest                   bd952fdf9b43        12 seconds ago      6.15GB
pytorch/pytorch         latest                   37b81722dadc        9 days ago          4.16GB

起動方法

/home/[ユーザー名]/ML/にデータを保存するとします。

$ docker run -d --rm --gpus all -v ~/ML:/workspace -p 8888:8888 --name pytorch pytorch-lab jupyter lab

オプション
-d バックグラウンドで実行
--rm コンテナ終了時にコンテナを削除
--gpus all コンテナでGPUを使用可能にする
-v ~/ML:/workspace ローカルの~/MLをコンテナの/workspaceにマウントして作成したファイルが保存されるようにする。
-p 8888:8888 ローカルのポート8888への接続をコンテナの8888(jupyterlab)へ送る。
--name pytorch コンテナに任意の名前をつける。

使い方

ブラウザからhttp://localhost:8888/に接続、ログインページが表示されるので設定したパスワードでログインする。

問題点

• pytorchのイメージもライブラリもバージョンを指定せずビルドの度に最新のものがインストールされるので、今後のアップデート次第では作ったイメージがうまく動作しない可能性がある。
• rootでjupyterlabを起動するのはセキュリティ的にどうなんだろう(SSHポートフォワードを前提にはしている)。
• tqdmの表示はうまくいくが、torch vision.datasets.MNISTなどのダウンロード時の進捗バーは正常に表示されない。