Cudaバージョンに応じたPyTorchをインストールする方法

何をしたいか

• pip install torch==1.7.1+cu110のような、pypiでホストされていないバージョンをダウンロードしたい
• 結論："-f"オプションで、ダウンロード先をpypiでないPyTorchのURLに指定すればいい

状況

• PyTorchをインストールした際にcuda11.0を使ってインストールするようOfficialに書いてあったので、別環境でも同じようにインストールしようとしたらできなかった
• 最新版の更新に伴ってcu110のような枝番がpypiから削除されることが原因でした。

どうすればいいか

• pypiでなく、PyTorchのURLからダウンロードするようにすればいい

自分の使いたいcudaのバージョンを確認

cuda_11.0を使っている

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Wed_Jul_22_19:09:09_PDT_2020
Cuda compilation tools, release 11.0, V11.0.221
Build cuda_11.0_bu.TC445_37.28845127_0

もし複数のcudaが入っていたら注意

# もしdebファイルで11.0のほか例えば10.2もインストールした場合、10.2に変更したければ
# /usr/local/cudaのシンボリックリンクを11.0に向くように上書きすればいい
$ ls -l /usr/local
lrwxrwxrwx  1 root root   10  3月 18 11:57 cuda -> cuda-11.0/
$ rm /usr/local/cuda && ln -s /usr/local/cuda-11.0 /usr/local/cuda
$ ls -l /usr/local
lrwxrwxrwx  1 root root   10  3月 18 11:57 cuda -> cuda-10.2/

自分の環境にあったwheelを探す

• ここから空気を読んで探す

cuda_11.0に対応した1.7.1がほしい場合

# 自分はlinuxでpython3.8.6なので「cu110/torch-1.7.1%2Bcu110-cp38-cp38-linux_x86_64.whl」を発見
$ pip install torch==1.7.1+cu110 -f https://download.pytorch.org/whl/torch_stable.html