将 uv 与 PyTorch 配合使用

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

注意

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

安装 PyTorch

从软件包管理的角度来看，PyTorch 有一些不常见的特性：

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

因此，所需的软件包配置会因您需要支持的平台和想要启用的加速器而异。

首先，考虑以下（默认）配置，这可以通过运行 uv init --python 3.12 然后运行 uv add torch torchvision 生成。

在这种情况下，PyTorch 将从 PyPI 安装，PyPI 为 Windows 和 macOS 提供仅 CPU 的 wheel 文件，并为 Linux 提供 GPU 加速的 wheel 文件（针对 CUDA 12.4）：

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
]

支持的 Python 版本

在撰写本文时，PyTorch 尚未发布适用于 Python 3.14 的 wheel 文件；因此，requires-python = ">=3.14" 的项目可能无法解析。请参阅兼容性矩阵。

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

使用 PyTorch 索引

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

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

仅 CPUCUDA 11.8CUDA 12.1CUDA 12.4ROCm6英特尔 GPU

[[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-cu121"
url = "https://download.pytorch.org/whl/cu121"
explicit = true

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

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

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

我们建议使用 explicit = true，以确保该索引仅用于 torch、torchvision 以及其他与 PyTorch 相关的包，而像 jinja2 这样的通用依赖项，仍应从默认索引（PyPI）获取。

接下来，更新 pyproject.toml，将 torch 和 torchvision 指向所需的索引：

仅 CPUCUDA 11.8CUDA 12.1CUDA 12.4ROCm6英特尔 GPU

[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-cu121", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu121", 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-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
  { index = "pytorch-cu124", 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 支持依赖于 `pytorch-triton-rocm`，它也应该从 PyTorch 索引安装
# （并包含在 `project.dependencies` 中）。
pytorch-triton-rocm = [
  { index = "pytorch-rocm", marker = "sys_platform == 'linux'" },
]

PyTorch 没有为 macOS 发布英特尔 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'" },
]
# 英特尔 GPU 支持依赖于 `pytorch-triton-xpu`，它也应该从 PyTorch 索引安装
# （并包含在 `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.12.0"
dependencies = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
]

[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.12.0"
dependencies = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
]

[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.12.0"
dependencies = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
  "pytorch-triton-rocm>=3.3.0 ; 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.3"
explicit = true

或者，对于英特尔 GPU 构建版本：

[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12.0"
dependencies = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
  "pytorch-triton-xpu>=3.3.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 的构建版本，通过用户提供的额外参数来切换选择（例如，uv sync --extra cpu 与 uv sync --extra cu124）。

通过 tool.uv.sources，你可以使用额外标记为每个启用的额外参数指定所需的索引。例如，以下配置在 uv sync --extra cpu 时将使用仅支持 CPU 的 PyTorch 版本，在 uv sync --extra cu124 时使用支持 CUDA 的构建版本：

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

[project.optional-dependencies]
cpu = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
]
cu128 = [
  "torch>=2.7.0",
  "torchvision>=0.22.0",
]

[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 加速构建版本，因此当启用 cu124 额外参数时，上述配置在 macOS 上安装将会失败。

`uv pip` 接口

上述示例主要聚焦于 uv 的项目接口（uv lock、uv sync、uv run 等），不过 PyTorch 也可以通过 uv pip 接口进行安装。

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 索引：

$ UV_TORCH_BACKEND=auto uv pip install torch

启用此功能后，uv 将查询已安装的 CUDA 驱动程序版本，并为所有相关软件包（如 torch、torchvision 等）使用最兼容的 PyTorch 索引。如果未找到此类 CUDA 驱动程序，uv 将回退到仅支持 CPU 的索引。对于 PyTorch 生态系统之外的任何软件包，uv 将继续遵循现有的索引配置。

要选择特定的后端（例如 cu126），可设置 --torch-backend=cu126（或 UV_TORCH_BACKEND=cu126）。

目前，--torch-backend 仅在 uv pip 接口中可用，并且仅支持检测 CUDA 驱动程序（与 ROCm 或英特尔 GPU 等其他加速器不同）。

由于 --torch-backend 是预览功能，应将其视为实验性的，不受 uv 的标准版本控制策略约束。在 uv 的未来版本中，--torch-backend 可能会更改或完全移除。