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の中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー