解決策

$ pip install torch==1.5.0+cu101 torchvision==0.6.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

で解決

環境

Ubuntu
CUDA 10.1
NVIDIA-SMI 418.39
Python 3.7.12

起きたこと

Deep Learning用PCでPyTorchのセットアップをしている時に、PyTorchがGPUを認識してくれなかった。

Anacondaで仮想環境を構築後、condaでPyTorchをインストールした。

$ conda create -n test_torch python=3.7

conda activateで作った環境に移動した後、こちらを参考にしてPyTorchをインストールした。

$ conda install pytorch torchvision cudatoolkit=10.1 -c pytorch

これで無事PyTorchがインストールされ、使えるようになったと思った。

test.py

import torch

x = torch.rand(5, 3)
print(x)

しかし、これではCUDAが使えなかった。

cuda.py

import torch

torch.cuda.is_available()  # False

試したことと解決策

GPUを認識しているか試すために、以下を実行した

gpu.py

import torch

torch.cuda.current_device()

すると、Torch not compiled with CUDA enabledというエラーが出た。
このエラーで検索すると、こちらが見つかった。どうやら、condaでインストールしたことが悪さしているようだった。改めてpipでPyTorchをインストールし直した。

$ pip install torch===1.5.0 torchvision===0.6.0 -f https://download.pytorch.org/whl/torch_stable.html

インストール後、もう一度torch.cuda.current_device()とすると、次は以下のようなエラーが出た。

AssertionError: 
The NVIDIA driver on your system is too old (found version 9010).
Please update your GPU driver by downloading and installing a new
version from the URL: http://www.nvidia.com/Download/index.aspx
Alternatively, go to: https://pytorch.org to install
a PyTorch version that has been compiled with your version
of the CUDA driver.

どうやらNVIDIAのドライバが古いらしい。
ドライバのアップデートをすればいいのだろうが、このPCではこの環境に依存した他の実行環境も動いているため、できるだけアップデートは行いたくない。

そこで、このドライバのバージョンに合ったPyTorchをインストールすることにした。

公式から過去のバージョンのインストールコマンドを探す。今回はCUDA10.1のため、v1.5.0を使用した。

# CUDA 10.1
$ pip install torch==1.5.0+cu101 torchvision==0.6.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

インストール後、もう一度以下を実行すると、きちんとGPUが読み込まれた。

cuda.py

import torch

print(torch.__version__)  # 1.5.0+cu101
torch.cuda.is_available()  # True
torch.cuda.current_device()  # 0

まとめ

CUDA10.1でNVIDIAドライバをアップデートせずにPyTorchのインストールを行うには、以下のコマンドが有効である。

# CUDA 10.1
$ pip install torch==1.5.0+cu101 torchvision==0.6.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

どうやら肝は+cu101の部分らしい。CUDAの10.2以降であれば+cu102のようにCUDAのバージョン指定をしなくても良いが、CUDA10.1やCUDA9.2の場合は、指定が必要なようである。