【Python】UVでGPU依存のPyTorchをインストールする
ノーコール森本🤙AI電話SaaS CTO動機
-
poetry→uvに乗り換えたため
-
GPUありの環境・なしのマルチプラットフォームに対応したいため。
コーディングする際はMacbook m1で開発することが多いのですが、デプロイする環境はLinux + CUDA環境です。環境によって柔軟にPytorchのバージョンやタイプを変更したいというのが主な動機です。
poetry→uvへの移行はこの記事が参考になりました。
実装
[tool.uv.sources] のurlでインストールするwheelのurlを指定し、 marker でプラットフォームを指定すると、下記のようにインストール先の Wheel を振り分けられます。
例えば、sys_platform == 'darwin'とmarkerを指定することで、MacOS向けのPyTorchがインストールされます。
[project]
name = "project-name"
version = "0.1.0"
requires-python = ">=3.10"
dependencies = []
[tool.uv.sources]
torch = [
# macOS ARM（Apple Silicon）用
{ url = "https://download.pytorch.org/whl/cpu/torch-2.1.1-cp310-none-macosx_11_0_arm64.whl", marker = "sys_platform == 'darwin'" },
# Linux aarch64 (ARM) 用
{ url = "https://download.pytorch.org/whl/cpu/torch-2.1.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
# Linux x86_64 + CUDA 用
{ url = "https://download.pytorch.org/whl/cu121/torch-2.1.1%2Bcu121-cp310-cp310-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
]
しかし、今回の例では厳密にgpuの有無によって、PyTorchのインストールを振り分けているわけではありません。あくまでも今回はMacOSでの開発とLinux + GPU環境での本番運用のため、これで問題ないと判断しました。もっと良い実装例あれば教えてください🙇♂️
ノーコール森本🤙AI電話SaaS CTOその他試したこと
[project.optional-dependencies] と [tool.uv.sources] の組み合わせを使う方法が uv の公式ドキュメントに載っています。たとえば次のような pyproject.toml のイメージになります。
[project]
name = "project-name"
version = "0.1.0"
description = ""
readme = "README.md"
requires-python = ">=3.10"
dependencies = []
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "gpu" },
],
]
[project.optional-dependencies]
cpu = [
"torch>=2.1.1", ## torch==X.X.Xと、バージョンをfixするとtorch==2.1.1+cu121のwheelを参照してしまい、エラーが発生する？
"torchvision>=0.16.1",
"torchaudio>=2.1.1",
]
gpu = [
"torch>=2.1.1",
"torchvision>=0.16.1",
"torchaudio>=2.1.1",
]
[tool.uv.sources]
torch = [
{ index = "torch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'"}, ## pytorchのcpu版はlinuxとWindows向けでなので、Darwin（macOS）では無印をインストールする。https://docs.astral.sh/uv/guides/integration/pytorch/#configuring-accelerators-with-optional-dependencies
{ index = "torch-cu121", extra = "gpu"},
]
torchvision = [
{ index = "torch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'"},
{ index = "torch-cu121", extra = "gpu"},
]
torchaudio = [
{ index = "torch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'"},
{ index = "torch-cu121", extra = "gpu"},
]
[[tool.uv.index]]
name = "torch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "torch-cu121"
url = "https://download.pytorch.org/whl/cu121"
explicit = true
実行コマンド例
- cpuの場合
uv sync --extra cpu
- gpuの場合
uv sync --extra gpu
この設定だと、「最初の sync は成功しているのに、uv sync --reinstall --extra cpu のように再インストールすると失敗する」ケースがある、という事象に遭遇しました。下記のようなエラーメッセージが出ます。
error: Distribution torch==2.1.1+cu121 @ registry+https://download.pytorch.org/whl/cu121 can't be installed because it doesn't have a source distribution or wheel for the current platform
また、追加でuv add, uv removeをすると、解決に非常に時間がかかってしまうケースもありました。これはLayerXのブログにもあるように、sourceやversionによってリソースを指定すると、不要な環境のwheelを全て取得してしまうことが原因だと考えれられます。