跳到内容

使用 uv 管理 PyTorch

PyTorch 生态系统是深度学习研究和开发的热门选择。你可以使用 uv 来管理不同 Python 版本和环境下的 PyTorch 项目及依赖项,甚至可以控制加速器的选择(例如:仅 CPU 版本与 CUDA 版本)。

注意

本指南中概述的某些功能需要 uv 0.5.3 或更高版本。我们建议在配置 PyTorch 之前先进行升级。

安装 PyTorch

从打包的角度来看,PyTorch 有一些不同寻常的特性。

  • 许多 PyTorch 的 wheel 文件托管在专用索引中,而不是 Python 包索引 (PyPI)。因此,安装 PyTorch 通常需要将项目配置为使用 PyTorch 索引。
  • PyTorch 为每种加速器(如:仅 CPU、CUDA)生成不同的构建版本。由于发布或安装时没有标准化的机制来指定这些加速器,PyTorch 将其编码在本地版本说明符中。因此,PyTorch 的版本号通常看起来像 2.5.1+cpu2.5.1+cu121 等。
  • 不同加速器的构建版本发布在不同的索引中。例如,+cpu 构建版本发布在 https://download.pytorch.org/whl/cpu,而 +cu121 构建版本发布在 https://download.pytorch.org/whl/cu121

因此,必要的打包配置将根据你需要支持的平台和想要启用的加速器而有所不同。

首先,考虑以下(默认)配置,这是通过运行 uv init --python 3.14 后接 uv add torch torchvision 生成的。

在这种情况下,PyTorch 将从 PyPI 安装,PyPI 为 Windows 和 macOS 提供仅 CPU 的 wheel 文件,为 Linux 提供 GPU 加速的 wheel 文件(截至 PyTorch 2.9.1,目标为 CUDA 12.8)。

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14"
dependencies = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
]

对于希望在 Windows 和 macOS 上使用 CPU 构建,并在 Linux 上使用 CUDA 构建的项目,这是一个有效的配置。但是,如果你需要支持不同的平台或加速器,则需要相应地配置项目。

使用 PyTorch 索引

在某些情况下,你可能希望在所有平台上使用特定的 PyTorch 变体。例如,你可能也想在 Linux 上使用仅 CPU 的构建版本。

在这种情况下,第一步是将相关的 PyTorch 索引添加到你的 pyproject.toml 中。

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu118"
url = "https://download.pytorch.org/whl/cu118"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
[[tool.uv.index]]
name = "pytorch-rocm"
url = "https://download.pytorch.org/whl/rocm6.4"
explicit = true
[[tool.uv.index]]
name = "pytorch-xpu"
url = "https://download.pytorch.org/whl/xpu"
explicit = true

我们建议使用 explicit = true,以确保该索引用于 torchtorchvision 和其他与 PyTorch 相关的包,而不是像 jinja2 这样的通用依赖项,这些依赖项应继续从默认索引 (PyPI) 获取。

接下来,更新 pyproject.toml 以将 torchtorchvision 指向所需的索引。

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu" },
]
torchvision = [
  { index = "pytorch-cpu" },
]

PyTorch 不发布适用于 macOS 的 CUDA 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 在 Linux 和 Windows 上使用 PyTorch 索引,但在 macOS 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-cu118", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu118", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

PyTorch 不发布适用于 macOS 的 CUDA 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 将 PyTorch 索引仅用于 Linux 和 Windows,在 macOS 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

PyTorch 不发布适用于 macOS 的 CUDA 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 将 PyTorch 索引仅用于 Linux 和 Windows,在 macOS 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

PyTorch 不发布适用于 macOS 的 CUDA 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 将 PyTorch 索引仅用于 Linux 和 Windows,在 macOS 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu130", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

PyTorch 不发布适用于 macOS 或 Windows 的 ROCm6 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 将 PyTorch 索引仅用于 Linux,在 macOS 和 Windows 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
torchvision = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
# ROCm6 support relies on `pytorch-triton-rocm`, which should also be installed from the PyTorch index
# (and included in `project.dependencies`).
pytorch-triton-rocm = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]

PyTorch 不发布适用于 macOS 的 Intel GPU 构建版本。因此,我们通过 sys_platform 进行限制,指示 uv 将 PyTorch 索引仅用于 Linux 和 Windows,在 macOS 上回退到 PyPI。

[tool.uv.sources]
torch = [
  { index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
# Intel GPU support relies on `pytorch-triton-xpu`, which should also be installed from the PyTorch index
# (and included in `project.dependencies`).
pytorch-triton-xpu = [
  { index = "pytorch-xpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]

作为完整的示例,以下项目将在所有平台上使用 PyTorch 的仅 CPU 构建版本。

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
]

[tool.uv.sources]
torch = [
    { index = "pytorch-cpu" },
]
torchvision = [
    { index = "pytorch-cpu" },
]

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

使用环境标记配置加速器

在某些情况下,你可能希望在一种环境(例如 macOS 和 Windows)中使用仅 CPU 构建,而在另一种环境(例如 Linux)中使用 CUDA 构建。

使用 tool.uv.sources,你可以使用环境标记为每个平台指定所需的索引。例如,以下配置将在 Linux 上使用 PyTorch 的 CUDA 构建,并在所有其他平台(如 macOS 和 Windows)上使用仅 CPU 构建。

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]
torchvision = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
  { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
]

[[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

类似地,以下配置将在 Linux 上使用 PyTorch 的 AMD GPU 构建,并在 Windows 和 macOS 上使用仅 CPU 构建(通过回退到 PyPI)。

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
  "pytorch-triton-rocm>=3.5.1 ; sys_platform == 'linux'",
]

[tool.uv.sources]
torch = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
torchvision = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]
pytorch-triton-rocm = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]

[[tool.uv.index]]
name = "pytorch-rocm"
url = "https://download.pytorch.org/whl/rocm6.4"
explicit = true

或者,对于 Intel GPU 构建:

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
  "pytorch-triton-xpu>=3.5.0 ; sys_platform == 'win32' or sys_platform == 'linux'",
]

[tool.uv.sources]
torch = [
  { index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]
torchvision = [
  { index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]
pytorch-triton-xpu = [
  { index = "pytorch-xpu", marker = "sys_platform == 'win32' or sys_platform == 'linux'" },
]

[[tool.uv.index]]
name = "pytorch-xpu"
url = "https://download.pytorch.org/whl/xpu"
explicit = true

使用可选依赖项配置加速器

在某些情况下,你可能希望在某些情况下使用仅 CPU 构建,而在其他情况下使用 CUDA 构建,并通过用户提供的 extra 进行切换(例如 uv sync --extra cpuuv sync --extra cu128)。

使用 tool.uv.sources,你可以使用 extra 标记为每个启用的 extra 指定所需的索引。例如,以下配置将在 uv sync --extra cpu 时使用 PyTorch 的仅 CPU 版本,并在 uv sync --extra cu128 时使用 CUDA 构建。

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.14.0"
dependencies = []

[project.optional-dependencies]
cpu = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
]
cu128 = [
  "torch>=2.9.1",
  "torchvision>=0.24.1",
]

[tool.uv]
conflicts = [
  [
    { extra = "cpu" },
    { extra = "cu128" },
  ],
]

[tool.uv.sources]
torch = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]
torchvision = [
  { index = "pytorch-cpu", extra = "cpu" },
  { index = "pytorch-cu128", extra = "cu128" },
]

[[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

注意

由于 macOS 上没有 GPU 加速的构建版本,上述配置在启用 cu128 extra 时将无法在 macOS 上安装。

uv pip 接口

虽然上述示例侧重于 uv 的项目接口(uv lockuv syncuv run 等),但也可以通过 uv pip 接口安装 PyTorch。

PyTorch 本身提供了一个专用接口来确定给定目标配置应运行的 pip 命令。例如,你可以通过以下方式在 Linux 上安装稳定的、仅 CPU 的 PyTorch:

$ pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

要使用 uv 采用相同的工作流程,请将 pip3 替换为 uv pip

$ uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

自动后端选择

uv 支持通过命令行参数 --torch-backend=auto(或环境变量 UV_TORCH_BACKEND=auto)自动选择合适的 PyTorch 索引,如下所示:

$ # With a command-line argument.
$ uv pip install torch --torch-backend=auto

$ # With an environment variable.
$ UV_TORCH_BACKEND=auto uv pip install torch

启用后,uv 将查询已安装的 CUDA 驱动程序、AMD GPU 版本和 Intel GPU 的存在情况,然后为所有相关包(如 torchtorchvision 等)使用最兼容的 PyTorch 索引。如果未找到此类 GPU,uv 将回退到仅 CPU 的索引。对于 PyTorch 生态系统之外的任何包,uv 将继续遵守现有的索引配置。

你也可以使用 --torch-backend=cu126(或 UV_TORCH_BACKEND=cu126)选择特定的后端(例如 CUDA 12.8)。

$ # With a command-line argument.
$ uv pip install torch torchvision --torch-backend=cu126

$ # With an environment variable.
$ UV_TORCH_BACKEND=cu128 uv pip install torch torchvision

目前,--torch-backend 仅在 uv pip 接口中可用。