DockerでPytorch & Jupyter Lab環境を簡単構築する #Docker

PytorchとJupyter Labの環境をDockerで構築します。

以前は機械学習の環境構築にAnacondaとか使ってたんですが、なんか余計なものが沢山付いてくるし、ターミナルに(base)って表示されるのが好きじゃないので脱condaしました！

環境

Ubuntu 20.04.6 LTS
Intel(R) Core(TM) i7-9700K
GeForce RTX 2080 Ti

Dockerをインストールする

まずDockerをインストールしましょう。オフィシャルの手順に従えば良いかと思います。https://docs.docker.com/engine/install/ubuntu/
すでにインストールしてある方は飛ばしてください。

NVIDIAドライバをインストールする

NVIDIAドライバがインストールされているか確認します。$ nvidia-smiを実行します。なにも表示されない場合はインストールが必要です。すでにインストールしてある方は飛ばしてください。

まずリポジトリを追加します

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

次に以下のコマンドでドライバのリストを表示しrecommendedと付いた推奨ドライバを確認します。

$ ubuntu-drivers devices

確認できたらインストールし再起動しましょう

$ sudo apt install nvidia-driver-○○
$ sudo reboot

$ nvidia-smiを再度実行し情報が表示されたらOKです。

$ nvidia-smi

output

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.125.06   Driver Version: 525.125.06   CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
                                  :

CUDA Version: ○○　と出ますがここに表示されているCUDAのバージョンは現在インストールされているCUDAのバージョンでは無く、ドライバが対応している最新のバージョンです。

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

Dockerコンテナ内からローカルのGPUを利用するためにNVIDIA Container Toolkitなるものをインストールします。詳しい説明は省略しますが、以下を参考にしてください。
https://github.com/NVIDIA/nvidia-docker/tree/master#quickstart

Dockerfileの作成

Dockerfileを作成します。まず
https://hub.docker.com/r/nvidia/cuda　
からベースイメージを探します。
base, runtime, develの3つがありますが、develで問題ないと思います。
次に必要なライブラリを記述します。
今回はこんな感じにしました。

Dockerfile

FROM nvidia/cuda:11.3.0-cudnn8-devel-ubuntu20.04

RUN apt-get update
RUN apt-get install -y python3-pip

RUN python3 -m pip install --upgrade pip \
&&  pip install --no-cache-dir \
    black==22.3.0 \
    jupyterlab==3.4.2 \
    jupyterlab_code_formatter==1.4.11 \
    lckr-jupyterlab-variableinspector==3.0.9 \
    jupyterlab_widgets==1.1.0 \
    ipywidgets==7.7.0 \
    import-ipynb==0.1.4

RUN python3 -m pip install --no-cache-dir \
    torch==1.11.0+cu113 \
    torchvision==0.12.0+cu113 \
    torchaudio==0.11.0+cu113 \
    --extra-index-url https://download.pytorch.org/whl/cu113

RUN apt-get autoremove -y &&\
    apt-get clean &&\
    rm -rf /usr/local/src/*

WORKDIR /work

Dockerfileからdocker Imageを作成する

先ほど作成したDockerfileを元にイメージをビルドしましょう。

$ docker image build -t [IMAGE NAME] .

コンテナを起動する

以下のコマンドでコンテナを起動します。

$ docker run \ #コンテナ起動(&生成)
-v $PWD/[host_dir]:/[container_dir] \ #ローカルで作成した内容とコンテナ内で作成した内容が同期する設定
-w /[container_dir] \ #ワーキングディレクトリの設定
-it \ #インタラクティブな状態に設定
--gpus all \ #コンテナから使うGPUを指定
--rm \ #コンテナ終了時にコンテナを削除
-p 8888:8888 \ #ポートの設定 [ホストのポート番号:コンテナのポート番号]
[IMAGE_NAME:TAG] \

jupyter-lab \ #JupyterLab起動
--no-browser \ #勝手にブラウザを開かない
--ip 0.0.0.0 \ #サービスを IP アドレス 0.0.0.0 で LISTEN させる
               #(この指定がないと、JupyterLab 起動時に例外が送出されて起動失敗します)
--allow-root \ #rootでの実行は非推奨のため、このオプションをつけないと起動に失敗します
--NotebookApp.token='' \

起動するとhttp://localhost:8888でJupyter Labが開くはずです。

Jupyter Labを開いたらLancherからTerminalを起動して次のコマンドを実行しコンテナ上の仮想環境からpytorchでGPUを認識できているか確認しましょう。

$ python3 -c "import torch; print(torch.__version__); print(torch.cuda.is_available()); print(torch.cuda.get_device_name())"

output

1.13.0+cu116
True
NVIDIA GeForce RTX 2080 Ti

このように出力されればOKです。快適な機械学習ライフを！