使用UV管理PyTorch项目
PyTorch是深度学习研究和开发的流行选择。可以使用uv管理PyTorch项目,包括不同Python版本依赖、管理环境、甚至加速器选择等。
安装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 安装,其中包含适用于 Windows 和 macOS 的仅 CPU wheel,以及适用于 Linux 的 GPU 加速wheel(针对 CUDA 12.6):
[project]
name = "project"
version = "0.1.0"
requires-python = ">=3.12"
dependencies = ["torch>=2.7.0","torchvision>=0.22.0",
]
对于希望在 Windows 和 macOS 上使用 CPU 版本、在 Linux 上使用支持 CUDA 的版本的项目来说,这是一种有效的配置。不过,如果需要支持其他平台或加速器,则需要对项目进行相应配置。
使用PyTorch索引
在某些情况下,你可能希望在所有平台上使用特定的 PyTorch 变体。例如,你可能也想在 Linux 上使用仅支持 CPU 的构建版本。
这种情况下,第一步是将相关的 PyTorch 索引添加到你的 pyproject.toml
中:
CPU ONLY
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
CUDA 11.8
[[tool.uv.index]]
name = "pytorch-cu118"
url = "https://download.pytorch.org/whl/cu118"
explicit = true
CUDA 12.6