😚
uvでPyTorchのCPU / CUDAバージョンを環境ごとに管理する

2024/08/22に公開
Python
PyTorch
tech
こんにちは。今回はPythonのパッケージマネージャのuvで、PyTorchのCPU/CUDAのバージョンをLinuxやmacOSなどの環境に応じて切り替えてインストールする方法を紹介します。

 はじめに先日リリースされたuvのv0.30.0で、Ryeで実装されていたようなPython自体の管理もできるようになりました。
https://astral.sh/blog/uv-unified-python-packaging
https://nikkie-ftnext.hatenablog.com/entry/uv-0.3.0-release-awesome-future-python-project-package-manager
https://dev.to/astrojuanlu/python-packaging-is-great-now-uv-is-all-you-need-4i2d
https://zenn.dev/turing_motors/articles/594fbef42a36ee
コマンド体系はRyeとほぼ同じで、以下のような感じで好きなバージョンのPython環境を構築することができます。
uv python install 3.11
uv python pin 3.10
poetryではMulti Platformを考慮した依存関係解決とパッケージインストールが可能ですが、依存関係の解決が遅く、Ryeでは依存関係の解決はRust実装のおかげで高速なものの、特定のマシンの環境に依存したrequirements.lockが作成されてしまい、複数の環境に跨ったプロジェクトの管理がしづらいという問題がありました。

そこでpoetryと同じようにMulti Platformを考慮可能なuvがRyeのPython管理機能を吸収してくれたことで、両方の利用者が問題なく移行できるのではないかと思います。
この機能追加でpoetryやRyeで管理していたプロジェクトを移行しようと思い、PyTorchのインストール方法を試行錯誤していました。
https://zenn.dev/mjun0812/scraps/671db64dc42ffa
そこで、poetryのenvironment markerと同じような方法で、環境によってインストールするPyTorchを切り替える方法を発見したので共有します。
https://zenn.dev/nakakiiro/articles/3cc5f4080a7a09

 成果物まず最初に完成したpyproject.tomlを以下に示します。
[project]
name = "new-uv"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
    "torch==2.4.0+cu118; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torch==2.4.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
    "torchvision==0.19.0+cu118; sys_platform == 'linux' and platform_machine == 'x86_64'",
    "torchvision==0.19.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
    "numpy<2.0.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.uv]
find-links = [
    "https://download.pytorch.org/whl/cu118/torch",
    "https://download.pytorch.org/whl/cu118/torchvision",
]
上記のpyproject.tomlは以下の環境で検証しています。
macOS 14.6 Sonoma, arm64, M1 MacBook Air 2020
Ubuntu 24.04, aarch64, Dockerコンテナ(上記のmacOS上)
Ubuntu 22.04, x86_64, Ryzen 5700+RTX3090の自作PC

 解説uvではPEP508に基づいたパッケージの依存関係の記述が可能です。
https://docs.astral.sh/uv/concepts/dependencies/#pep-508
この記述方法の中にEnvironment Markerという機能があり、システムの環境ごとに依存関係を記述することができます。

上記のpyproject.tomlの該当箇所は以下です。
"torch==2.4.0+cu118; sys_platform == 'linux' and platform_machine == 'x86_64'",
"torch==2.4.0; sys_platform == 'darwin' or (sys_platform == 'linux' and platform_machine == 'aarch64')",
この記述では、
Linux + x86_64アーキテクチャの時はCUDA 11.8対応のPyTorch
macOS + arm64(Apple Silicon Mac)とLinux + aarch64の時はcpu版のPyTorch

がインストールされます。
また、PyTorchのバージョン表記2.4.0+cu118の+cu118の部分はLocal Versionと呼ばれており、uvでLocal Versionを使う時はパッケージのバージョン指定に==しか使えないことに注意して下さい。(>=は使えない)
https://docs.astral.sh/uv/pip/compatibility/#local-version-identifiers
そして、以下でも言及されていますが、macOSでtorchをインストールする場合はtorch==2.4.0+cpuではなく、torch==2.4.0と記述しないとインストールできません。

これは今後改善されるかもしれないです。
https://github.com/astral-sh/uv/issues/5182
以上、uvでのPyTorchのインストール方法の紹介でした。uvはPythonの管理機能が追加されて、隙が無くなった印象です。今後注目していきたいと思います。
はじめに

成果物

解説

Discussion
