PyTorchをバイナリでインストールするとCUDAがついてくる

GPU搭載マシンが欲しかったため、年末セールでゲーミングPCを購入しました。

購入したのはAcer Nitro N50 のGeForce RTX 3060搭載モデルです。

機械学習関係のプログラムを実行するためにCUDAをインストールしようと思ったのですが、その途中でPyTorchのバイナリでインストールするとCUDAが自動的に入ることに気がついたので記録しておきます。

参考: The PyTorch binaries include the CUDA and cuDNN libraries. · Amy Tabb

検証手順

1: まずはUbuntu 22.04.1 LTSをインストールしました。インストール時の設定は以下の通りです。

最小インストール
インストール中にアップデートをダウンロードしない
サードパーティ製ソフトウェアをインストールしない
Windowsと併用可能な形でインストール

2: UbuntuにPython環境をインストールしました。

sudo apt install python3-dev python3-pip

3: PyTorchのGET STARTEDの指示に従って、PyTorchをpipでインストールしました。

pip install torch torchvision torchaudio

このとき、nvidia-cuda-runtime-cu11 や nvidia-cudnn-cu11 などのパッケージが同時にインストールされました。

4: CUDAが有効になっているかどうか以下のようにPythonを対話モードで実行して確認しました。

$ 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
>>> torch.cuda.is_available()
False
>>> torch.cuda.device_count()
0
>>> torch.cuda.current_device()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/kivantium/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 552, in current_device
    _lazy_init()
  File "/home/kivantium/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 229, in _lazy_init
    torch._C._cuda_init()
RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

どうやらNVIDIAのdriverは別途有効にしないといけないようです。

5: Ubuntuの「追加のドライバー (Additional Drivers)」を立ち上げると、複数のドライバーバージョンが表示されました。

バージョンを確認するためにエラーメッセージで示されたURLからNVIDIA Driver Downloadsに飛び、自分の環境に合ったドライバーのバージョンを確認しました。

v525.60が指定されていたので、nvidia-driver-525を選んで適用した後、マシンを再起動しました。

6: 再起動後に再びPythonを立ち上げ、CUDAが有効化されていることを確認しました。

$ 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
>>> torch.cuda.is_available()
True
>>> torch.backends.cudnn.version()
8500
>>> torch.cuda.device_count()
1
>>> torch.cuda.current_device()
0

というわけで、PyTorchしか使わないのであればCUDAを別途インストールする必要はなさそうです。