安装

按照官方指引,可以全局安装,也可以在具体的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())"