Ubuntu22.04でPyTorchを使うためのCUDA+cuDNNの環境構築

UbuntuでGPUを使ってPyTorchの計算がしたい。

この記事は以下を参考にわかりやすいようにまとめた。
[PyTorchとGPU/CUDA周りの環境構築のバージョン解決](https://qiita.com/ketaro-m/items/4de2bd3101bcb6a6b668)
[[Linux] ubuntu20.04のPCにCUDA, cudnnをインストールする](https://qiita.com/2yuu/items/e62367007380e3c564f7)

## 環境
- Ubuntu 22.04
- GeForce RTX3070

## モチベーション
PyTorchでDNNの学習を行うとき、1epochの学習時間がWindowsよりもUbuntuの方が3倍速かったので、Ubuntu環境が欲しい。

## Windowsだと遅い原因
Windows用のNVIDIA GPUドライバにはTCC(Tesla Compute Cluster)とWDDM(Windows Display Driver Model)という2種が存在している。
- TCCドライバはCUDAカーネルの起動のオーバーヘッドを削減するため、深層学習などの複雑な計算に適しているが、Teslaシリーズのグラフィックを無効にする。
- WDDMドライバはWindowsでのグラフィックスアクセラレーションに使用される。

### 遅くなる原因
WDDMモードではcudaで使用されるGPUメモリを90%に制限される。
また、cudaプロセスはcudaプロセス自身に90%のGPUメモリしか使用できない。
つまり1つのプロセスで 90% × 90% = 81% のGPUメモリしか使用できないようになっている。
（他にも原因がありそうだがこれしか見つけられなかった。）

よってTCCドライバを使いたいが、このTCCドライバは**GeForceシリーズに対応していない**。
Ubuntu(Linux)では複数のドライバモデルはサポートされていないため（WDDMがいらないから？）、UbuntuでのGPUの方が早くなるっぽい。

・GPUメモリが制限される件 [Pytorch is slower on windows than on linux #22083](https://github.com/pytorch/pytorch/issues/22083#issuecomment-552741570)
・TCCについて [4. Tesla Compute Cluster (TCC)](https://docs.nvidia.com/nsight-visual-studio-edition/reference/index.html#tesla-compute-cluster)
・Driver Modelについて（pp.8-9） [nvidia-smi.txt](https://developer.download.nvidia.com/compute/DCGM/docs/nvidia-smi-367.38.pdf)

というわけでUbuntuの環境構築をする。

## 環境構築の流れ
1. NVIDIAドライバのインストール
2. CUDAのインストール
3. cuDNNのインストール
4. PyTorchのインストール

## 1. NVIDIAドライバのインストール
[ここから](https://developer.nvidia.com/cuda-downloads) 環境にあったものを選択。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1294427/b00fdf78-c506-4c03-3c3c-be8f29bed46b.png)
下部に出てくる**Driver Installer**から、**To install the open kernel module flavor:** と書かれた下側2行を実行する。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1294427/6d4024f3-449b-0084-db2a-32060154b91d.png)


```
$ sudo apt-get install -y nvidia-kernel-open-545
$ sudo apt-get install -y cuda-drivers-545
```


再起動して`nvidia-smi`を実行し、GPUが認識されているか確認する。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1294427/d019a5f7-a203-8713-a932-1aa811845a96.png)


このような表示が出ていれば完了。

>右上に`CUDA Version: 12.3`などと表示されるが、インストールされているCUDAバージョンではなく、**互換性のある最新のCUDAバージョン**を示している。そのためCUDAが未インストールでも表示される。
[引用元](https://qiita.com/ketaro-m/items/4de2bd3101bcb6a6b668)


## 2. CUDAのインストール
[先ほどのリンク](https://developer.nvidia.com/cuda-downloads) で生成された上側の**Base Installer** を実行する。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1294427/ff582c74-0d75-153f-7c45-0b92f502b7b5.png)

```
$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
$ sudo dpkg -i cuda-keyring_1.1-1_all.deb
$ sudo apt-get update
$ sudo apt-get -y install cuda-toolkit-12-3
```

インストールが完了したら再起動して、`~/.bashrc`を開き、
```
$ sudo nano ~/.bashrc
```
末尾に以下の2行を追記し、`Ctrl + O`で保存、`Ctrl + X`で閉じる。
```
export PATH=/usr/local/cuda:/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/lib:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
```
その後、以下のコマンドを実行する。
```
$ source ~/.bashrc
```
最後に`nvcc -V`で以下のような表示が出れば完了。
```
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Fri_Sep__8_19:17:24_PDT_2023
Cuda compilation tools, release 12.3, V12.3.52
Build cuda_12.3.r12.3/compiler.33281558_0
```

## 3. cuDNNのインストール
まずZlibをインストール
```
$ sudo apt-get install zlib1g
```

次に、[ここから](https://developer.nvidia.com/rdp/cudnn-archive) CUDAのバージョンに合わせたcuDNNのバージョンを選択し、
`Local Installer for Ubuntu22.04 x86_64 (Deb)` をダウンロードする。
なお、ダウンロードにはNVIDIAのアカウント作成が必要になる。 
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1294427/bece5a02-da0f-b2cf-9f2f-9ff28e8dd014.png)

その後、公式[1.3.2. Debian Local Installation](https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html#installlinux-deb)に従いインストールを行う。
はじめに以下のコマンドを実行。`x`はダウンロードしたファイルに合わせる。
```
$ cd ~/Download
$ sudo dpkg -i cudnn-local-repo-ubuntu2204-8.x.x.x_1.0-1_amd64.deb
$ sudo cp /var/cudnn-local-repo-*/cudnn-local-*-keyring.gpg /usr/share/keyrings/
$ sudo apt update
```
その後以下のコマンドを実行する。`x`はダウンロードしたファイルに合わせ、
**`X.Y`はCUDAのバージョン`-1`を入れる。**
例：インストールしたのが`CUDA12.3` なら、...+cuda`X.Y`は`cuda12.2`にする。
[補足2](https://qiita.com/Manyan3/items/628b3b22700fd569e8fb)
```
$ sudo apt-get install libcudnn8=8.x.x.x-1+cudaX.Y
$ sudo apt-get install libcudnn8-dev=8.x.x.x-1+cudaX.Y
$ sudo apt-get install libcudnn8-samples=8.x.x.x-1+cudaX.Y
```

## PyTorchのインストール
ここまで来たらあと少し。
仮想環境はお好みで用意してもらい、`torch` & `torchvision`をインストールする。
```
$ pip3 install torch torchvision
```
以下のPythonコードを実行し、
```test.py
import torch
import torchvision

# バージョン確認
print("Python      :", platform.python_version())
print("torch       :", torch.__version__)
print("torchvision :", torchvision.__version__)

# GPUの確認
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device      :", device)
```
`device`が`cuda`と表示されていればすべて完了。
```出力例
Python      : 3.10.12
torch       : 2.1.0+cu121
torchvision : 0.16.0+cu121
device      : cuda
```

## 以上、お疲れ様でした。