Windows11+WSL2でGPUのPyTorchを最低限動かすためのオレオレDockerfile

---

Windows11＋WSL2でGPU対応のPyTorchのDockerイメージを1から作る方法を紹介します。CUDAとCuDNNはDockerhubで公開されているnvidiaのイメージを用い、ライブラリのインストールは各自行います。あなただけの最強環境を作ろう！

pip環境のDockerイメージがほしい

とりあえず使うだけならこちらで紹介したnvidia-dockerで十分ですが、nvidia-dockerのPythonはconda環境なので、pipでのインストールとの衝突が気になります。そこで、オレオレDockerfileでGPUのPyTorch環境を作成してしまう方法を考えました。

まずはそのDockerfileから。

# PyTorchのバージョンに合わせてCUDAのバージョンを変える
FROM nvidia/cuda:11.3.0-runtime-ubuntu20.04
WORKDIR /workspace
RUN apt-get update
RUN apt-get install -y software-properties-common tzdata
ENV TZ=Asia/Tokyo 
RUN add-apt-repository ppa:deadsnakes/ppa
# pipの他に最低限使いそうなもの
RUN apt-get install -y python3-pip git vim wget
# Ubuntuデフォルトのpython3をpythonに置き換えるようにシンボリックリンクの作成
RUN ln -s /usr/bin/python3.8 /usr/bin/python
# PyTorchのバージョンにあわせて変える
RUN pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
RUN apt-get upgrade -y

これを「Dockerfile」という名前で保存し（拡張子なし）、Linuxのカレントディレクトリに配置します。

元のイメージはDockerhubにあるnvidia cudaから。これは前述のnvidia-dockerと異なり、CUDAとCuDNNが入っているだけのイメージで、PyTorchすら入っていません。カスタムして自分でいろいろ突っ込みたいときは便利そうです。

ビルドはWSL2がインストールされている前提で、スタート→「Ubuntu on Windws」を開き、以下のようにします。

sudo docker build -t pytorch .

これは「pytorch」というタグをつけたイメージを作ることをさします。その他各自必要なライブラリがある場合はDockerfileに追記しましょう。

Dockerを起動するときはnvidia-dockerと同じで、

sudo docker run --gpus all -it --rm --shm-size=4g pytorch

shm-sizeはお好みで。

Dockerイメージのサイズ

これでビルドしたときのイメージサイズは以下の通りです。

REPOSITORY    TAG                          IMAGE ID       CREATED         SIZE
pytorch       latest                       1c4cfb9a3ac2   9 minutes ago   7.89GB
nvidia/cuda   11.3.0-runtime-ubuntu20.04   0e0c38a8d6e0   2 months ago    1.77GB

8GB弱ぐらいです。nvidia-dockerとそんなに変わらなそうです。

WSL2だから訓練に失敗するケースもある

WSL2は自分で書いたPyTorchのソースぐらいは普通に動くのですが、論文サンプルのように他人の書いた複雑なソースは必ずしも動かないことがあります。

このようにNCCL errorが出ることがあります。とある論文のサンプルコードからです。以下のことを試してみましたがエラーは消えませんでした。

• docker runで「–shm-size=32g」のオプションをつけてメモリ領域を増やす
• コード内でlocalhostにアクセスしている部分があるので、「localhostForwarding=True」の設定ファイルをWindowsで作成（参考）
• Distribute Trainingに対応していない可能性があるので、GPU1個で訓練

このエラーが手強かったので最終的に諦めてしまいました。いくらスクラッチビルドしても、ネイティブのLinux環境とはまだ完全な互換性はないようです。ハマったら諦めたほうが精神衛生的にはいいかもしれません。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---