poetry + python 複数バージョン + pytorch の環境構築

この記事は？

[motivation] poetry で python3.7系＋python3.8系を許容しつつ、pytorch も入れたい
[conclusion] 技術的には可能だが実用的には無理

背景

• pytorch を使ったライブラリを作成している
  • pip install でインストールされて利用されることを想定している
• ライブラリは python3.7 と python3.8 をサポートする必要がある

メイン

poetry 単体でもできることはできるが実行時間が異様に長く UX が悪い

poetry の issue に上がっている pytorch のインストール方法 を参考に書くとこうなります。

pyproject.toml

[tool.poetry.dependencies]
python = ">=3.7,<3.9"  # python3.7 & 3.8 を許容
torch = { version = "1.9.1", source = "torch"}
torchvision = { version = "0.10.1", source = "torch"}

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cu111"  # cuda version に合わせてリンクは適当に変える
secondary = true

確かにこれでもできるっちゃできるんですがこの方法で poetry lock を実行すると cu111 の wheel ファイルをマシンのアーキテクチャや指定した python version にかかわらず全てダウンロードします。上記の pyproject.toml を lock すると以下の画像のハイライト部分の wheel ファイルを全てダウンロードしてから依存を解決しにいくため、めちゃくちゃ時間がかかり¹ UX が異常に悪く、頻繁に poetry lock を実行する場合現実的ではないです。

代替手段は？

cpu 版 pytorch を poetry で管理することで依存を解消し、gpu 版の pytorch を後から入れ直すと速いです。

pyproject.toml

[tool.poetry.dependencies]
python = ">=3.7,<3.9"
torch = "==1.9.1"
torchvision = "==0.10.1"

poetry install
pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html

結論

• poetry lock がめちゃくちゃかかってもいい場合は source を指定して依存を全て解かせると安心。コードも綺麗。ただし本当に時間がかかる
• poetry lock に時間をかけたくない場合は pip を併用していい感じにやる

備考

• 扱う python version が3つ以上でも多分同じことができます
• wheel を全部落とすのをやめてください

1. 自分は比較的高速な回線でやっていますがそれでも 6分/wheel 程度かかります。github の issue では 250Mbps で 15-20分/wheel かかったと報告されています。 ↩