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

uv v0.4.23でindex-urlを複数指定してインストールする方法が追加されました。 以下の新しい記事で紹介しています。 この記事で紹介している方法だと、CUDA 12.4以降は動作しない場合があるため、 以下の新しい記事を参考にしてください。

uvでPyTorchのCPU/CUDAバージョンを環境ごとに管理する その2 | MJUN Tech Note

とある情報系の院生の技術ノートです

https://mjunya.com/posts/2024-10-22-python-uv-pytorch-2

はじめに

先日リリースされたuvのv0.30.0で、Ryeで実装されていたようなPython自体の管理もできるようになりました。

uv: Unified Python packaging

Manage entire Python projects and even Python itself with a single unified tool.

https://astral.sh/blog/uv-unified-python-packaging

uv 0.3.0リリース、Pythonプロジェクト管理ツールのすがた！ uv、お前、Ryeを屠るのか... - nikkie-ftnextの日記

はじめに 黄前久美子部長、お誕生日おめでとうございます！！ nikkieです1 今週のPyCoder's Weekly (#643)を読んでいたところ、めちゃ驚く記事が飛び込んできました。 この記事を読んでぶち上がったテンションでの走り書きです。 uv導入のための検証兼素振りは別途やりたいなと思っています。 目次 はじめに 目次 uv、パッケージインストーラのすがた uv、パッケージマネージャのすがた Python開発に必要なすべてをまとめてきている！ 終わりに uv、パッケージインストーラのすがた uvは「高速なpip」として認識していました。 Rustで実装されているので速いという理解です…

https://nikkie-ftnext.hatenablog.com/entry/uv-0.3.0-release-awesome-future-python-project-package-manager

コマンド体系は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のインストール方法を試行錯誤していました。

uvでPyTorchをインストールする方法の試行

https://zenn.dev/mjun0812/scraps/671db64dc42ffa

そこで、poetryのenvironment markerと同じような方法で、環境によってインストールするPyTorchを切り替える方法を発見したので共有します。

tool.poetry.sourceを使わずにPyTorchを高速にpoetry管理下におく

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'",
    "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 22.04, x86_64, Ryzen 5700+RTX3090の自作PC

解説

uvではPEP508に基づいたパッケージの依存関係の記述が可能です。

Dependencies | uv

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)の時はcpu版のPyTorch がインストールされます。

また、PyTorchのバージョン表記2.4.0+cu118の+cu118の部分はLocal Versionと呼ばれており、uvでLocal Versionを使う時はパッケージのバージョン指定に==しか使えないことに注意して下さい。(>=は使えない)

Compatibility with pip | uv

https://docs.astral.sh/uv/pip/compatibility/#local-version-identifiers

そして、以下でも言及されていますが、macOSでtorchをインストールする場合はtorch==2.4.0+cpuではなく、torch==2.4.0と記述しないとインストールできません。 これは今後改善されるかもしれないです。

`uv add` rejects pytorch from its index when on mac · Issue #5182 · astral-sh/uv

I am trying to use uv to add torch to my project's dependencies, but it fails. I am on macOS14.4.1 (ARM). Here is a minimal repro. [project] name = "test" version = "0.1.0" requires-python = ">=3.1...

https://github.com/astral-sh/uv/issues/5182

以上、uvでのPyTorchのインストール方法の紹介でした。uvはPythonの管理機能が追加されて、隙が無くなった印象です。今後注目していきたいと思います。