PoetryでGPU版Pytorchをクイックインストール

PoetryでサクッとGPU版Pytorch環境を作成したいとき用の備忘録
Poetry1.5.1から以下の方法が有効です

1. 以下のサイトからインストールコマンドを確認します

https://pytorch.org/get-started/locally/
https://pytorch.org/get-started/previous-versions/

今回は以下の例を使用します

pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu121

2. --priority=explicitオプション をつけて、インストールコマンドのindex-urlをpoetry source add コマンドで登録します

poetry source add torch_cu121 --priority=explicit https://download.pytorch.org/whl/cu121

※ torch_cu121は任意の名前で大丈夫です

3. sourceを指定して、poetry add でインストールする

poetry add torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --source torch_cu121

Poetryは指定したsource(torch_cu121)を探しに行き、環境に沿ったwheelファイルを見つけ、インストールしてくれるようです

4. pip listでGPU版Pytorchがインストールされていることを確認します

...

torch              2.1.1+cu121
torchvision        0.16.1+cu121
torchaudio         2.1.1+cu121

※ pyproject.tomlは以下のようになります

project.toml

...

[tool.poetry.dependencies]
python = "^3.10"
torch = {version = "2.1.1", source = "torch_cu121"}
torchvision = {version = "0.16.1", source = "torch_cu121"}
torchaudio = {version = "2.1.1", source = "torch_cu121"}

[[tool.poetry.source]]
name = "torch_cu121"
url = "https://download.pytorch.org/whl/cu121"
priority = "explicit"

...