対応するCUDAバージョンがないときのPytorchのソースからのビルド

はじめに

機械学習をしているとよくPytorchを使用する機会がありますが, インストール済みのCUDAのバージョンに対応するPytorchのバイナリファイルがないことがあります. そういった場合, Pytorchをソースからビルドし, CUDAに対応するPytorchをインストールすることができます. 今回はPytorchをソースからビルドした例を紹介します.

セットアップ

仮想環境を作成しておきます.

python3 -m venv torch
source torch/bin/activate

CUDAとPytorchのバージョンが合うか確認します.
CUDA11.4に合うPytorchのバージョンは1.10.0のようです.

CUDAバージョンを指定してPytorchのソースを取得します.

# get sources
git clone https://github.com/pytorch/pytorch.git -b v1.10.0

依存パッケージをインストールしておきます. Pytorchのビルド中もパッケージ不足で中断することがあるので, その都度インストールします.

# install dependency
pip install numpy
pip install pyyaml
pip install typing_extensions

Pytorchのビルドをします.

cd pytorch
python setup.py install

ImportError: Failed to load PyTorch C extensions:が出る場合

python setup.py develop && python -c "import torch"

PytorchでGPUを使用できるか以下のコマンドで確認します. Trueが返ってこれば, GPUが使用可能になっています.

# launch python terminal
python

# import Pytorch
import torch

# command to confirm GPU available or not 
print(torch.cuda.is_available())
# True -> available GPU

対応するtorchvisionもインストールしておきます. torchのバージョン1.10.0に対応するtorchvisionのバージョンは0.11.1です.

git clone https://github.com/pytorch/vision.git -b v0.11.1
cd vision
python setup.py install

おわりに

いつもはビルド済みのバイナリファイルでPytorchをインストールすることが多いですが, 今回はソースからビルドしてみました. ぜひお試しを!