WSL2のUbuntuにCUDAをインストールしてPyTorchを動かす

タイトル通りですが、WSL2のUbuntuにCUDAをインストールして、PyTorchでCUDAを利用する方法を紹介します。

そこまで難しい設定はありませんが、2021年4月17日時点ではWindows Insider Programに参加して、Windows10のビルドバージョンを20145以上にする必要があります。

多少面倒なのとdev版のビルドバージョンを利用する必要がありますが、WSLでGPUを使えるようになるのは便利すぎます。

おそらくインストール方法は随時変わります。下記の方法を試す前にNVIDIAの公式チュートリアル(CUDA on WSL)の確認をお勧めします！

目次非表示

実行環境
Windows10のビルドバージョンを20145以上へ
Windows10にNVIDIA driverをインストール
WSLにCUDA Toolkitをインストール
WSL環境にPyTorchをインストール
まとめ

実行環境

Windows10 Pro (OSビルド 21359)
Ubuntu 20.04 on WSL2
GeForce MX350
CUDA Toolkit 11.3
Python 3.8
PyTorch 1.8.1

さらっと書きましたが、GeForce MX350でCUDA&PyTorchを使おうとする人ってどのくらいいるのでしょうか、、、

Windows10のビルドバージョンを20145以上へ

NVIDIAの公式チュートリアルで言及されていますが、Windows10のビルドバージョンを20145以上にする必要があります（推奨は21332以上）。

WindowsBlogsによると、20150以上でWSLでのGPU計算が対応したようです。（バージョンが異なる理由は謎、、、）

現在のOSのバージョンは”winver”で確認できますよ。スタートボタンを押して、”winver”で検索してみましょう。

この記事を書いた時（2021/04/17）では、20145以上のビルドバージョンにするためには、Windows Insider Programに登録してDevチャンネルからWindowsアップデートを行う必要があります。手順は下記の通りです。

公式サイトからWindows Insider Programに登録
システム設定から「Windows Insider Programの設定」を開き、Insiderの設定を選択して下さいの箇所でDevチャンネルを選択
システム設定から「Windows Update」を開きアップデートのチェックをして、更新
再起動

Insider Programへの登録などは特に難しいことはないので、案内に従って行っていきましょう。

Windows10にNVIDIA driverをインストール

NVIDIAの公式ページ（CUDA on Windows Subsystem for Linux (WSL) – Public Preview）からdriverをダウンロードして、インストールしましょう。GeForceかQuadroかを選択する必要があるので、自分の環境に合った方を選択します。

また、driverのダウンロードにはNVIDIA Developerへの登録が必要になります。

WSLにCUDA Toolkitをインストール

WSL上のUbuntuにCUDA Toolkitをインストールしていきます。NVIDIAがWSL用にインストールスクリプトを公開しているので、それを利用しましょう（CUDA Toolkit Downloads）。

下記に自分が利用したコマンドを載せておきますが、バージョンなどは注意してください。自分はインストール時の最新版である11.3をインストールしました。

この記事の冒頭で紹介したCUDA on WSLのページでも言及されていますが、WSLにnvidiaのdriverはインストールしてはいけません。Windows側にインストールしたnvidia driverにすでにWSL用のdriverが含まれているからです。

cudaをインストールできたら、下記でWSLがGPUを使えているかを確認してみましょう。

$ cd /usr/local/cuda/samples/4_Finance/BlackScholes
$ sudo make
$ ./BlackScholes
[./BlackScholes] - Starting...
GPU Device 0: "Pascal" with compute capability 6.1
...

GPU Device...というログが出力されていれば、GPUが認識されています。

WSL環境にPyTorchをインストール

ここではvenvで仮想環境を作成して、PyTorchをインストールしてみます。バージョンは各自適切なものを入れましょう。

また、PyTorchのインストールコマンド(pip)はPyTorchの公式サイトから手に入れてください。

# 仮想環境の作成
$ python3 -m venv ~/envs/pytorch-cuda
$ source ~/envs/pytorch-cuda/bin/activate
# PyTorchのインストール
$ pip install "torch==1.8.1+cu111" "torchvision==0.9.1+cu111" "torchaudio==0.8.1" -f https://download.pytorch.org/whl/torch_stable.html
# gpuが利用可能であることを確認
$ python -c "import torch; print(torch.cuda.is_available())"
True