コンテンツにスキップ

uv を使った torch の依存管理

uv とは

uv は 2024 年 2 月に公開された非常に高速な Python パッケージインストーラー兼リゾルバー。pip と pip-compile を置き換えるものとして Rust で書かれており、公式が出しているベンチマークによるとpippip-toolsより 10 倍から 100 倍速い。 静的バイナリとして提供されているため、Python や Rust のインストールが不要である。 依存解決の戦略として highest 以外にも lowest が選択できたり、依存のオーバーライドも可能。

この記事の目的

uv はまだ開発段階でありpippip-toolsと比べて機能が不足している部分もある。 これまでできていたことが現時点の uv だとできないこともあり、その一つが torch のインストールである。 uv v0.1.10 でも torch 単体のインストールはできるが、torchvision のように torch に依存するパッケージと一緒にインストール場合は工夫が必要だったり、これまでのpipと同じようなパッケージの指定方法だとエラーが出ることもある。

このような状況でも少しの工夫で torch の依存解決とインストールができるので、その方法を紹介する。

uv のインストールと使い方

インストール

# On macOS and Linux.
curl -LsSf https://astral.sh/uv/install.sh | sh

# On Windows.
powershell -c "irm https://astral.sh/uv/install.ps1 | iex"

# With pip.
pip install uv

# With pipx.
pipx install uv

# With Homebrew.
brew install uv

# With Pacman.
pacman -S uv

ファイルからロック済みのパッケージを出力

uv pip compile pyproject.toml -o requirements.txt   # Read a pyproject.toml file.
uv pip compile requirements.in -o requirements.txt  # Read a requirements.in file.

uv を使って torch の依存解決

先述の通り、uv で torch の依存解決とインストールを行うには少し工夫が必要である。 pip-compile で依存解決を行い、その出力を pip でインストールという作業をそのまま uv で実行するとエラーが起こる。 この作業をエラーなく uv で行うためには、overrideの機能を使うと良い。

override はその単語の通り、依存のバージョンを上書きすることができる。

1
2
3
# overrides.txt
torch==2.1.1+cpu
torchvision==0.16.1+cpu

上記のようなファイルを作成し、uv pip compile--overrideオプションで指定することで、torch の依存解決とインストールが可能になる。

uv pip compile pyproject.toml -o requirements.txt --override overrides.txt
uv pip install -e . --override overrides.txt

参考