WSL2のUbuntu上にPyTorchとGPU環境構築

PyTorch

2022.12.302023.03.28

今回は、WSL2のUbuntu上にPyTorchとGPU環境構築の方法についてまとめてみました。
※注：CUDAのバージョンをPyTorchに合わせてCUDA11.7になるように修正しました。

前準備
GPU関係アプリのインストール
終わりに
関連記事

前準備

前準備として、以下の環境を準備します。
・Windows11のWSL2環境
・Ubuntu22.04
・Python3.10.6(Ubuntu22.04のデフォルトのまま)
なお、PCの構成は以下の通りです。

PCの構成(参考)

OS：Windows11HOME(22H2)
CPU：12thGen Intel Core i5-12400F (2.5GHz)
GPU: NVIDIA GeForce RTX 3080(10GB)
RAM: 32GB

VisualStadioCodeのUbuntu対応

エディタはVisualStadioCodeを使うとファイルの編集が楽にできます。あらかじめ、公式ページよりVisualStadioCodeをダウンロードして、インストールします。Ubuntu用の拡張機能もあるので、つけておきます。

一旦、VisualStadioCodeを閉じて、Ubuntuのターミナルを開いて「code .」とすることで、Ubuntu環境でVisualStadioCodeが開きます。

code .

GPU関係アプリのインストール

NVIDIAドライバの更新

下記のNVIDIAのWebサイトから、構成にあった最新のドライバを確認し、最新になっていない場合は、念のため、インストールします。

最新の NVIDIA 公式ドライバーをダウンロード

NVIDIA GPUサポートのセットアップ

ここからはNVIDIAのサポートページの指示に従ってVSCodeのUbuntuのターミナルに入ってWSL2 でCUDAとNVIDIA GPUサポートをセットアップします。下記のようにすることで、WSL2にCUDA11.7をインストールできます。

$ sudo apt-get update
$ sudo apt-key del 7fa2af80
$ wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
$ sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
$ wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb
$ sudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb
$ sudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
$ sudo apt-get update
$ sudo apt-get -y install cuda
$ /usr/local/cuda/bin/nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Tue_May__3_18:49:52_PDT_2022
Cuda compilation tools, release 11.7, V11.7.64
Build cuda_11.7.r11.7/compiler.31294372_0

cuDNNのインストール

ブラウザでnvidiaのテクニカルサポートからUbuntu22.04とcudaのバージョンを確認し、cuDNNをダウンロードします。ダウンロードしたdebファイルはubuntuのホームディレクトリにコピーします。

cuDNN Archive

Download releases from the GPU-accelerated primitive library for deep neural networks.

ダウンロードしたdebファイルでcuDNNをインストールします。

$ sudo dpkg -i cudnn-local-repo-ubuntu2204-8.6.0.163_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2204-8.6.0.163/cudnn-local-FAED14DD-keyring.gpg /usr/share/keyrings/

下が2004版

$ sudo dpkg -i cudnn-local-repo-ubuntu2004-8.6.0.163_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-ubuntu2004-8.6.0.163/cudnn-local-B0FE0A41-keyring.gpg /usr/share/keyrings/

PyTorchのインストール

引き続いてPyTorchをインストールします。公式ページから、インストールコマンドを確認し、そのまま入力します。

$ sudo apt -y update
$ sudo apt -y install python3-dev python3-pip python3-setuptools
$ pip install -U pip setuptools
$ source .profile
$ pip3 install torch torchvision torchaudio

PyTorchの起動確認

PyTorchでGPUが使えるか確認してみます。torchによるGPUの認知とGPUを使ったテンソルの定義ができることが分かります。

$ python3
Python 3.10.6 (main, Nov 14 2022, 16:10:14) [GCC 11.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.__version__)
1.13.1+cu117
>>> print(torch.cuda.is_available())
True
>>> print(torch.cuda.get_device_name())
NVIDIA GeForce RTX 3080
>>> device = torch.device('cuda:0')
>>> a = torch.zeros([2, 4], dtype=torch.int32, device=device)
>>> print(a)
tensor([[0, 0, 0, 0],
        [0, 0, 0, 0]], device='cuda:0', dtype=torch.int32)
>>> exit()

終わりに

WSL2のUbuntuにPyTorchとGPUの環境構築をまとめました。以前、ubuntu+condaで使ったことがあったのですが、個人的にubuntu中でのconda環境は構造がうまく理解できなかったので、シンプルにubuntuに直に置いた方が理解しやすかったです。。。ちょっとやりたいことがあるので、今回の環境をベースに少し調整してようと思います。