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

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 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アップデートを行う必要があります。手順は下記の通りです。

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

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

まとめ

WSL2のUbuntu20.04にCUDAをインストールしてPyTorchを動かすところまでを紹介しました。

Windows Insider Programに入ってdevチャンネルのビルドバージョンを利用しないといけないので、多少面倒です。ただ、手順はそこまで難しくないと思います。

CUDA関連はバージョン管理が非常に面倒なので、自分はいつもDockerにまとめています。WSL上でGPUが使えるようになったということは、WSL上でのDockerでもGPU使えるのでしょうか、、、このあたりも調べてみようと思っています。

WSLは本当に便利ですよね。Windowsの安定版ビルドにWSLでGPUを使える機能が早く入ることを期待しています。