環境

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
・TCCについて 4. Tesla Compute Cluster (TCC)
・Driver Modelについて（pp.8-9） nvidia-smi.txt

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

環境構築の流れ

NVIDIAドライバのインストール
CUDAのインストール
cuDNNのインストール
PyTorchのインストール

1. NVIDIAドライバのインストール

ここから環境にあったものを選択。

下部に出てくるDriver Installerから、To install the open kernel module flavor: と書かれた下側2行を実行する。

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

再起動してnvidia-smiを実行し、GPUが認識されているか確認する。

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

右上にCUDA Version: 12.3などと表示されるが、インストールされているCUDAバージョンではなく、互換性のある最新のCUDAバージョンを示している。そのためCUDAが未インストールでも表示される。
引用元

2. CUDAのインストール

先ほどのリンクで生成された上側のBase Installer を実行する。

$ 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

次に、ここから CUDAのバージョンに合わせたcuDNNのバージョンを選択し、
Local Installer for Ubuntu22.04 x86_64 (Deb) をダウンロードする。
なお、ダウンロードにはNVIDIAのアカウント作成が必要になる。

その後、公式1.3.2. Debian Local Installationに従いインストールを行う。
はじめに以下のコマンドを実行。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 なら、...+cudaX.Yはcuda12.2にする。
補足2

$ 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

以上、お疲れ様でした。