WSL2 上の PyTorch に GPU を認識させて深層学習環境をつくる

本記事では、WSL2 上の PyTorch から GPU を使用できる深層学習用の環境を構築していきます。環境は下記の通りです。

• OS : Windows 11 Home (21H2)
• CPU : AMD Ryzen 7 3700X 8 コア
• RAM : 16 GB
• GPU : NVIDIA GeForce RTX 2080 Ti

作業フロー

GPU を利用して PyTorch で記述したモデルの学習をするまでの道のりはこんな感じ。

1. Windows に NVIDIA ドライバーをインストール
2. WSL (Ubuntu) に CUDA Toolkit をインストール
3. WSL (Ubuntu) に cuDNN をインストール
4. WSL (Ubuntu) から GPU を認識できているか動作確認
5. WSL (Ubuntu) に PyTorch をインストール

用語辞典

NVDIA ドライバー

• NVIDIA 製のグラフィックボードを動かすために必要なドライバー
• 自分の使っているグラフィックボードに合ったドライバーを使う必要がある

CUDA (Toolkit)

• 本来、GPUは画像処理用の演算装置ですが、汎用的な数値計算を GPU 上で実行できるように NVIDIA が開発したプログラミング環境 (コンパイラやライブラリ、API などが含まれる)

cuDNN

• DNN で使われる基本的な機能をまとめた CUDA ライブラリ
• フレームワークごとに CUDA をコードを書くムダを無くしている
• NVIDIA 自身が最適化しているので性能的にも良いものを手軽に作れる

Windows に NVIDIA ドライバーをインストール

公式 NVIDIA ページ からドライバーをダウンロードしてインストールします。筆者のマシンに載っている GPU は、 NVIDIA GeForce RTX 2080 Ti なので下記のように入力してドライバーを検索しました。

Download Type という項目が、初見だと何のことかわからないと思いますが、下記のような考え方で良いと思います。

• Game Ready Drivers : とにかく発売されたばかりの最新のゲームがプレイできるようにしたい場合
• Studio Drivers : 安定性や品質を重視する場合

CUDA Toolkit のインストール

WSL 上の Ubuntu に CUDA Toolkit をインストールしましょう。公式からインストールまでのスクリプトが提供されている ので、これを叩いていくだけです。

cuDNN のインストール

まずは Windows 上で、公式ページ から deb ファイルをダウンロードします。

ぼくの環境は Ubuntu 20.04 で CPU は Ryzen 7 だったので Local Installer for Ubuntu20.04 x86_64 (Deb) をダウンロードしました。

Windows の C ドライブは、WSL 上の /mnt/c にマウントされているので、ここからダウンロードフォルダを参照して cuDNN をインストールします。

sudo dpkg -i /mnt/c/Users/utahk/Downloads/cudnn-local-repo-ubuntu2004-8.5.0.96_1.0-1_amd64.deb

WSL から GPU を認識できているか確認

~ $ nvidia-smi
Fri Aug 12 11:31:21 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 516.94       CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| 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 ...  On   | 00000000:09:00.0  On |                  N/A |
| 35%   29C    P8    22W / 250W |   1023MiB / 11264MiB |      2%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

~ $ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Jun__8_16:49:14_PDT_2022
Cuda compilation tools, release 11.7, V11.7.99
Build cuda_11.7.r11.7/compiler.31442593_0

PyTorch のインストール

公式 PyTorch サイト で、自身の環境にあったインストールコマンドを調べることができます。

あとは、これを WSL 側で叩いて PyTorch をインストールします。(下記のワンライナーは筆者の環境の場合のものなので注意)

pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

PyTorch のインストールには少々時間がかかると思いますが、インストールが終わったら PyTorch のユーティリティメソッドで PyTorch から CUDA を利用可能か確認することができますので、動作確認をしておきましょう。

import torch

torch.cuda.is_available()

これで、True が返ってくれば OK です。

まとめ

久しぶりにセットアップしましたが、Windows での環境設定が数年前より格段に楽になっているような印象を受けました。(気のせいかも)

また、普段は Mac を使っているので気づいていませんでしたが、WSL の導入がとても快適で Ubuntu が使えるようになるまでの体験がスムーズ過ぎて感動的でした。