【Python】PyTorchのGPU版を使うのに少しつまづいた

【状況】PyTorchで画像処理系の重い処理は厳しい．PyTorchにはGPU版があるけど，標準ではCPU版がインストールされるので，自分でインストールする必要がありました．
【対処】公式サイトに行ってダウンロード．ただし，使っているGPU（CUDAのバージョン）などによって選ぶ必要があります．

公式サイトを見る

Start Locally

執筆時点で以下のような選択肢があります．なお，Python 3.8以降が必要，とあります．コマンドをコピーし，コマンドプロンプトにコピーするとダウンロードされます．

ダウンロード内容を決める画面．クリックして環境を決めるとコマンドが得られる

やっかいなのが，CUDAの最新版には対応していないことです．執筆時点で，CUDA12.3が最新ですが，PyTorch2.1.2の対応はCUDA11.8，12.1となっています．
さらに，GPUが古くてCUDA11.7以前の場合は，古いPyTorchを使う必要があります．別のページにまとめられているので，該当するバージョンを探します．

Previous PyTorch Versions

インストール準備

必要であれば，CUDAをダウンロードします．12.1なら以下です．普通にたどると12.3になってしまいますね．Windows11用をダウンロードすると3.1GBとなっています．

CUDA Toolkit 12.1 Downloads

CUDA11.7, 11.8が入っていましたが，すんなりインストールできました．

インストール

PyTorch 2.1.2，Windows，Pip，Python，CUDA12.1の場合は，以下のコマンドになりました．

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

torchが2.5GB．なかなか大きいですね．
ちなみに，過去にいろいろなCUDAのバージョンを試したときに，更新できない場合がありました．そのときは，アンインストールしてからインストールしました．

pip uninstall torch torchvision torchaudio

インストール成功の旨のメッセージを見るとバージョン番号の後ろに，「+cu???」とCUDAのバージョン番号が入っていることが確認できます（ん？torchvisionがない？？）．

Successfully installed torch-2.1.2+cu121 torchaudio-2.1.2+cu121

念のために「pip list」を入力して確認（torchvisionにcuが無い）．

torch              2.1.2+cu121
torchaudio         2.1.2+cu121
torchvision        0.16.2

実行時にエラーが出て再インストール

とりあえず，PyTorch CPU版で動作していたプログラムを実行したところ，torchvisionのエラーが大量に表示されて動作せず・・・．一旦，すべてアンインストールしました．

pip uninstall torch torchvision torchaudio

そして，上記と同一のダウンロードを実行しました．

Successfully installed torch-2.1.2+cu121 torchaudio-2.1.2+cu121 torchvision-0.16.2+cu121

あ．torchvisionにcu付いてる！！「pip list」を見ると．

torch              2.1.2+cu121
torchaudio         2.1.2+cu121
torchvision        0.16.2+cu121

めでたくGPU版で動きました♪