跳转到内容

在PyTorch中使用uv

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

注意

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

安装 PyTorch

从打包的角度来看，PyTorch 有一些不常见的特性：

  • 许多PyTorch的wheel包托管在专用索引上，而非Python Package Index (PyI)。因此，安装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安装，该平台为Windows和macOS提供仅CPU的wheel包，并为Linux提供针对CUDA 12.4优化的GPU加速wheel包：

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

[[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指向所需的索引：

[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 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.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上回退到PyPI仅使用CPU版本：

[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

或者，对于Intel 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，您可以使用额外的标记来为每个启用的附加项指定所需的索引。例如，以下配置将使用 PyTorch 的仅 CPU 版本执行 uv sync --extra cpu，而使用支持 CUDA 的构建执行 uv sync --extra cu124：

[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上安装。

The uv pip interface

虽然上述示例主要聚焦于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

自动后端选择

在preview版本中，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 或 Intel GPU）。

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