UbuntuでGPUを使ってPyTorchの計算がしたい。
この記事は以下を参考にわかりやすいようにまとめた。
PyTorchとGPU/CUDA周りの環境構築のバージョン解決
[Linux] ubuntu20.04のPCにCUDA, cudnnをインストールする
環境
- 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コードを実行し、
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