Windows環境向けにWSLとDockerを用いてローカルでのPytorch環境を構築します。ローカルでサクッと試したいことができた際に、毎回クラウドコンピューティング環境を立ち上げてみたいなことをするのがめんどくさかったので、手元にあるGPUを活用してPytorchの環境を構築してみたいと思います。GPU向けのPytorch環境構築のため、cudaを使って環境を整えたいと思います。
想定環境
以下の環境を想定してインストールしていきます。
- OS: Windows11
- GPU: NVIDIA GeForce RTX 3060
- WSL2 - Ubuntu22.04
WSL2をインストール
まずはWSL2をインストールします。Linuxサブシステムを有効化して、好きな環境をインストールしましょう。Ubuntuのインストールについては、以下の記事で備忘録として解説しています。詳細な方法は、こちらをご確認ください。
Dockerをインストール
WSLのインストールが完了したらターミナルを起動して、Dockerのインストールを行いましょう。WSLでインストールしたOSにDockerを入れる場合は基本的に無償で使えます。過去の記事でDockerのインストールについて解説していますので、こちらをご確認ください。
cudaをインストール
GPUを使いたいので、ドライバであるcudaをインストールしていきます。Windows11またはWindows10の21H1だと基本的にサポートされているようなので、必要に応じてそれらにアップデートください。以下のMSのドキュメントが参考になります。
MS 公式Docs: https://learn.microsoft.com/ja-jp/windows/ai/directml/gpu-cuda-in-wsl
すでにそのバージョンになっている方は上記はスキップして、以下のNVIDIAのDocsからcudaのインストールをしてもらえればよいかと思います。
NVIDIA 公式 Docs: https://docs.nvidia.com/cuda/wsl-user-guide/index.html#getting-started-with-cuda-on-wsl-2
インストールした後にnvidia-smiと打ってみて、GPUやcudaのバージョンが表示されれば問題なくインストールできていると思います。上側には現在のGPUの使用状況、下側には実行中のプロセスが表示されます。
hoge@hoge:/mnt/c/Users/hoge/Desktop$ nvidia-smi
Sun Oct 1 XX:XX:XX 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.XX Driver Version: 552.XX CUDA Version: 12.2 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 3060 On | 00000000:01:00.0 On | N/A |
| 30% 29C P8 15W / 200W | 0MiB / 8192MiB | 0% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
+-----------------------------------------------------------------------------------------+Dockerfileを作成してPytorch環境を作る
環境構築が一通り完了したので、DockerFileを作成してしていきます。
ベースイメージを探す
まず自分の環境のcudaに合うdocker imageのベースを探します。わたくしの環境では12.2以下で構築する必要があるため、cudaが12.1のubuntuのベースイメージを探します。探すときは以下のリンクにあるdocker hubのnvidiaのレポジトリから探せばOKです。
docker hub : https://hub.docker.com/r/nvidia/cuda/tags
検索ボックスがあるので、そこにバージョンを打ち込んで下にスクロールしていけば、ubuntuのベースイメージがあるのでそれを元にpytorch環境を構築していきます。
以下のイメージを元にpytorchを入れていきます。
pytorchのinstallを行う
pytorchのインストールを行います。ただここでは実際に行うわけではなく、pytorchをinstallするコマンドをDockerfileに書くため、そのコマンドを調べに行きます。pytorchの以下の公式リンクから調べることが可能です。
以下のようなGUIが出てくるので自分の条件に合うものを選択して赤丸のところに出てくるインストールコマンドを表示させます。今回自分の場合は、pipで入れる想定で、以下セッティングとしました。
ここでベースイメージのcudaのバージョンと合わせる必要がある可能性があります。可能であれば、合わせておくとのちのトラブルが減るかもしれません。今回はcuda12.1に統一して構築していきます。
Dockerfileを作成
実際に上記で調べたことを元にDockerfileを作成していきます。以下はあくまで最小限だと思いますので、お好みで必要なものを追加してもらえればと思います。
FROM nvidia/cuda:12.1.0-base-ubuntu20.04
RUN apt-get update && apt-get install -y --no-install-recommends \
python3.10 \
python3-wheel \
python3-setuptools \
python3-pip \
curl \
vim \
wget \
git \
git-lfs \
build-essential \
pkg-config \
&& pip install --upgrade pip && pip3 install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu121 \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# Update symbolic link
RUN ln -s /usr/bin/python3 /usr/bin/python
RUN rm /usr/bin/pip
RUN ln -s /usr/bin/pip3 /usr/bin/pip
Docker build & run
実際に上記で作成したDockerfileを実際にbuildします。
docker build -t pytorch-learn:latest -f docker/Dockerfile .buildに成功したら以下のコマンドでrunしてみましょう。立ち上がって、nvidia-smiでGPUが見えていれば成功です。
docker run -it --rm --gpus all pytorch-learn:latest /bin/bash
root@2a3f66ed8753:/# nvidia-smi
Sun Oct 1 XX:XX:XX 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.XX Driver Version: 552.XX CUDA Version: 12.2 |
|-----------------------------------------+------------------------+----------------------+
| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+========================+======================|
| 0 NVIDIA GeForce RTX 3060 On | 00000000:01:00.0 On | N/A |
| 30% 28C P8 16W / 200W | 803MiB / 8192MiB | 4% Default |
| | | N/A |
+-----------------------------------------+------------------------+----------------------+
+-----------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=========================================================================================|
| 0 N/A N/A 52 G /Xwayland N/A |
+-----------------------------------------------------------------------------------------+エラーTips
もし以下のエラーが出る場合には、nvidia-container-runtimeを設定しなおすとうまくいきます。
docker: Error response from daemon: could not select device driver "" with capabilities: [[gpu]].以下のコマンドでupdateして、docker engineを再起動したのちに、runコマンドを行うことで起動すると思います。
sudo apt-get install nvidia-container-runtime
curl -s -L https://nvidia.github.io/nvidia-container-runtime/gpgkey | sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-container-runtime/$distribution/nvidia-container-runtime.list | sudo tee /etc/apt/sources.list.d/nvidia-container-runtime.list
sudo apt-get update
sudo apt-get install nvidia-container-runtime
sudo service docker restartまとめ
無事上記の流れでWSL2+docker+pytorchの環境を構築することができました。今回はGPUを使いたかったのでcudaを入れましたが、cpuでも動作させることができると思いますので、お好みでカスタマイズしてもらえればと思います。