【Python】poetry でCPU/GPU用のtorch 環境を用意する

はじめに

何度もpoetry で　pytorch の環境を構築しようとして失敗していました。poetry add すると cuda 用のpytorch や torchvision がインストールされてしまい、import torchvison すると、cucaについて xxxxxx.so がない、と言われて終了、という場面に遭遇しました。

ここではpoetry のCUIコマンドのみで(pip3無しで)環境を用意できたのでメモしておきます。CPU用のwheel ファイルを持ってくるように指定しますが、pyproject.toml を編集しません。

また、GPU用の場合は、とりあえず直接wheelファイルをpyproject.toml に記入して環境設定しています。もっとスマートにしたいですが、とりあえずということでメモっておきます。(2023/08/28追記)

内容

Python3.9でCPU用のtorch をpoetry でインストール

poetry init --python=3.9

で初期化した後、ダウンロードするサイトを指定します。

$ poetry source add -p explicit pytorch https://download.pytorch.org/whl/cpu
Adding source with name pytorch.

すると、pyproject.yaml は下記のように書き換えられました。

pyproject.yaml

[tool.poetry.dependencies]
python = "3.9.17"

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

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"

そしてインストールします。

poetry add --source pytorch torch torchvision

pyproject.toml には以下のように書かれていました。

[tool.poetry.dependencies]
python = "3.9.17"
torch = {version = "^2.0.1+cpu", source = "pytorch"}
torchvision = {version = "^0.15.2+cpu", source = "pytorch"}


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

下記のように無事に動きました。ほ。何時間も費やしました。とほほ。

$ poetry run python -V
Python 3.9.17
$ poetry run python -c "import torch; print(torch.__version__)"
2.0.1+cpu
$ poetry run python -c "import torchvision; print(torchvision.__version__)"
0.15.2+cpu

はまりポイント

最初 python 3.10 で実行していたのですが、どうにもこれだと torchvision をimport したときに、lzmaをimport できない、、、というエラーが出ます。これは標準ライブラリのようです。いろいろ調べると、ライブラリをインストールしてpython をbuild しなおせ、3.7, 3.10で起きる現象、のようなことが書かれていました。もう、、、と思い、3.9にして使うことにしました。

念のため、下記を実行しました。

￥sudo apt-get install -y build-essential libssl-dev zlib1g-dev libbz2-dev \
libreadline-dev libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev \
xz-utils tk-dev libffi-dev liblzma-dev python-openssl git

python 3.9.17の用意

pyenv でインストールしたのち、poetry のenv に指定しました。

pyenv は下記のようにしました。適切なのか良く分かりませんが。

$ pyenv install 3.9
$ pyenv local 3.9.17
$ pyenv versions
  system
  3.8.6
  3.8.16
* 3.9.17 (set by /home/x77/thistest2/.python-version)
  3.10.6

poetry の仮想環境で指定するようにします。

$ poetry env use 3.9.17

Python3.9でGPU用のtorch をpoetry でインストール(2023/08/28追記)

WSL2でのCUDA環境の用意はこちらに書いています。今回も新しくCUDA11.8を入れました。

あまり頭が良くないですが、poetry init したあと、pyproject.toml に直接入れたいものを書きました。

pyproject.toml

[tool.poetry.dependencies]
(略)
torch = {url =
"https://download.pytorch.org/whl/cu118/torch-2.0.0%2Bcu118-cp39-cp39-linux_x86_64.whl"}

下記の記事にしっかり書かれています。ここにあるようにpackage soruce を指定してwheel ファイルを明示的に書かない方式にしたかったのですが、source の参照でエラーになってしまったので今回は直接wheel ファイルを指定しました。

もっと粘ればできたかも。うーん。

まとめ

とりあえずこれでtorch, torchvision を使ったコードをCPUで動かせそうです。良かった。また少し長く生き延びられそうだ。
(2023/08/06)

GPU用の部分を追記しました。package source 指定にしたかったのですが確認の時間がなく。。。(2023/08/28)