この記事では、古いグラフィックボードでPyTorchを使うためにPyTorchをビルドする方法を記載します。

なぜPyTorchのビルドが必要か

私のPCはUbuntu 20.04で、グラフィックボードはGeForce 710 (Compute Capability 3.5)です。
最新のPyTorch(1.10)をpipでPyTorchをインストールしてみます。

pip3 install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html

しかし、pipでインストールしたPyTorchでは、Compute Capability 3.7以上のGPUにしか対応していません。

print(torch.cuda.get_arch_list())
['sm_37', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86']

pipでインストールしたPyTorchのCompute Capabilityの対応表を見ると、1.8.0まではCompute Capability 3.5に対応していることが分かりますが、1.8.0が対応しているCUDAのバージョンは11.1です。
そこで、CUDA11.1のインストールを試みますが、現時点ではasm/kmap_types.hが無いエラーでインストールできません。現時点のUbuntu 20.04のカーネルではインストールできなくなっているようです。
そこで、新しいCUDAを使って自分でPyTorch 1.10をビルドして使うことにします。

CUDAのインストール

CUDA 11.4以上はKepler対応はdeprecatedになっているようですので、11.3をインストールしました。
現時点ではsudo apt install nvidia-dkms-465でエラーが出ますが、このパッチをあてると解決しました。

PyTorchのビルド

whlを作ります。ビルドにはPythonが必要なので、システムのPythonを使います。

sudo apt install python3 python3-pip python3-numpy
pip3 install typing-extensions
git clone https://github.com/pytorch/pytorch --branch v1.10.2 --depth 1 pytorch-v1.10.2
cd pytorch-v1.10.2
TORCH_CUDA_ARCH_LIST="3.5" MAX_JOBS=2 time python3 setup.py bdist_wheel

distディレクトリにwhlができますので、PyTorchをインストールしたい環境でpip install xxx.whlしてインストールします。

pip install pytorch-1.10.2/dist/torch-1.10.0a0+git71f889c-cp38-cp38-linux_x86_64.whl

なお、ビルド時の変数を確認したいときは、TORCH_CUDA_ARCH_LIST="3.5" python3 setup.py build --cmake-onlyです。この場合、sudo apt install cmake-curses-guiしてccmake buildでビルドできます。
cleanしたいときは、python3 setup.py cleanです。

インストールしたら、Compute Capability 3.5に対応していることを確認します。

$ python
Python 3.8.2 (default, Jul 18 2021, 19:32:07)
[GCC 9.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.cuda.get_arch_list())
['sm_35']

Torchvisionのビルド

PyTorchをインストールした環境で以下を実行します。

git clone https://github.com/pytorch/vision --branch v0.11.1 --depth 1 vision-v0.11.1
cd vision-v0.11.1
MAX_JOBS=2 python3 setup.py install