uv对于gpu torch 的支持与管理
安装
按照官方指引,可以全局安装,也可以在具体的pyhon解释器用pip安装(换句话说可以和全局解释器共存,比如类似miniconda的版本管理器)
https://github.com/astral-sh/uv
背景
最近在研究pytorch的gpu版本安装, 早期使用的时候只能使用
uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
没有和uv集成,导致每次安装都要手动指定index-url,非常不方便
解决
最近发现uv对于pytorch的gpu版本支持的很好,可以参考如下链接
https://docs.astral.sh/uv/guides/integration/pytorch/
实战
准备pyproject.yaml文件
[project]
name = "test-project"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"annoy>=1.17.3",
"faiss-cpu>=1.11.0",
"gensim>=4.3.3",
"mysql-connector-python==8.0.28",
"pandas>=2.3.0",
"peewee>=3.18.1",
"redis>=6.2.0",
"ruff>=0.11.12",
"scikit-learn>=1.6.1",
]
[project.optional-dependencies]
cpu = [
"torch>=2.7.0",
"torchvision>=0.22.0",
]
cu128 = [
"torch>=2.7.0",
"torchvision>=0.22.0",
]
cu124 = [
"torch>=2.6.0",
"torchvision>=0.21.0",
]
[tool.uv]
conflicts = [
[
{ extra = "cpu" },
{ extra = "cu128" },
{ extra = "cu124" },
],
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", extra = "cu128" },
{ index = "pytorch-cu124", extra = "cu124" },
]
torchvision = [
{ index = "pytorch-cpu", extra = "cpu" },
{ index = "pytorch-cu128", extra = "cu128" },
{ index = "pytorch-cu124", extra = "cu124" },
]
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true
安装
uv add torch --extra cu124
验证
uv run python -c "import torch; print(torch.cuda.is_available())"
- 原文作者:大鱼
- 原文链接:https://brucedone.com/archives/1704/
- 版权声明:本作品采用知识共享署名-非商业性使用-禁止演绎 4.0 国际许可协议. 进行许可,非商业转载请注明出处(作者,原文链接),商业转载请联系作者获得授权。