OA0 = Omni AI 0
OA0 是一个探索 AI 的论坛
现在注册
已注册用户请  登录
OA0  ›  代码  ›  TorchInductor — PyTorch 编译后端

TorchInductor — PyTorch 编译后端

 
  compliance ·  2026-03-03 00:02:27 · 9 次点击  · 0 条评论  

PyTorch Logo


PyTorch 是一个 Python 包,主要提供两大高级功能:
* 支持 GPU 加速的张量计算(类似 NumPy)
* 基于自动求导系统的深度神经网络

你可以根据需要,复用你喜欢的 Python 包(如 NumPy、SciPy 和 Cython)来扩展 PyTorch。

我们的主干健康状况(持续集成信号)可在 hud.pytorch.org 查看。

关于 PyTorch 的更多信息

学习 PyTorch 基础知识

从细节层面看,PyTorch 是一个由以下组件构成的库:

组件 描述
torch 一个类似 NumPy 的张量库,提供强大的 GPU 支持
torch.autograd 一个基于磁带的自动微分库,支持 torch 中所有可微分的张量操作
torch.jit 一个编译栈(TorchScript),用于从 PyTorch 代码创建可序列化和可优化的模型
torch.nn 一个与 autograd 深度集成的神经网络库,旨在实现最大的灵活性
torch.multiprocessing Python 多进程,但能在进程间神奇地共享 torch 张量的内存。对数据加载和 Hogwild 训练很有用
torch.utils 为方便起见提供的 DataLoader 和其他实用函数

通常,PyTorch 被用作:
* 替代 NumPy,以利用 GPU 的强大算力。
* 一个提供最大灵活性和速度的深度学习研究平台。

进一步阐述:

支持 GPU 的张量库

如果你用过 NumPy,那么你就用过张量(又称 ndarray)。

张量示意图

PyTorch 提供的张量可以运行在 CPU 或 GPU 上,并能大幅加速计算。

我们提供了多种张量例程来加速并满足你的科学计算需求,例如切片、索引、数学运算、线性代数、归约操作。而且它们非常快!

动态神经网络:基于磁带的自动求导

PyTorch 有一种独特的构建神经网络的方式:使用并回放磁带记录器。

大多数框架(如 TensorFlow、Theano、Caffe 和 CNTK)对世界的看法是静态的。你必须构建一个神经网络并反复复用相同的结构。改变网络的行为方式意味着需要从头开始。

在 PyTorch 中,我们使用一种称为反向模式自动微分的技术,它允许你零延迟或零开销地任意改变网络的行为方式。我们的灵感来自该主题的几篇研究论文,以及当前和过去的工作,例如 torch-autogradautogradChainer 等。

虽然这项技术并非 PyTorch 独有,但它是迄今为止最快的实现之一。你可以为你的疯狂研究获得最佳的速度和灵活性。

动态图

Python 优先

PyTorch 不是将 Python 绑定到一个庞大的 C++ 框架上。它是为了深度集成到 Python 中而构建的。你可以像使用 NumPy/SciPy/scikit-learn 等一样自然地使用它。你可以使用你喜欢的库,在 Python 本身中编写新的神经网络层,并使用诸如 CythonNumba 这样的包。我们的目标是在适当的地方不重复造轮子。

命令式体验

PyTorch 的设计目标是直观、思路线性且易于使用。当你执行一行代码时,它就会被执行。不存在对世界的异步视图。当你进入调试器或收到错误信息和堆栈跟踪时,理解它们是直截了当的。堆栈跟踪会精确指向你代码定义的位置。我们希望你再也不必因为糟糕的堆栈跟踪或异步且不透明的执行引擎而花费数小时调试代码。

快速且轻量

PyTorch 的框架开销极小。我们集成了加速库,如 Intel MKL 和 NVIDIA (cuDNNNCCL),以最大化速度。其核心的 CPU 和 GPU 张量及神经网络后端已经成熟并经过多年测试。

因此,无论你运行小型还是大型神经网络,PyTorch 都非常快。

与 Torch 或其他一些替代方案相比,PyTorch 的内存使用效率极高。我们为 GPU 编写了自定义内存分配器,以确保你的深度学习模型具有最高的内存效率。这使你能够训练比以前更大的深度学习模型。

轻松扩展

编写新的神经网络模块,或与 PyTorch 的张量 API 对接,被设计得直截了当且抽象最少。

你可以使用 torch API 或你喜欢的基于 NumPy 的库(如 SciPy) 在 Python 中编写新的神经网络层。

如果你想用 C/C++ 编写你的层,我们提供了一个方便的扩展 API,它高效且样板代码最少。无需编写包装代码。你可以查看此教程此示例

安装

二进制安装包

通过 Conda 或 pip wheel 安装二进制文件的命令在我们的网站上:https://pytorch.org/get-started/locally/

NVIDIA Jetson 平台

NVIDIA Jetson Nano、Jetson TX1/TX2、Jetson Xavier NX/AGX 和 Jetson AGX Orin 的 Python wheel 包在此处提供,L4T 容器在此处发布。

它们需要 JetPack 4.2 及以上版本,由 @dusty-nv@ptrblck 维护。

从源码安装

前提条件

如果从源码安装,你需要:
* Python 3.10 或更高版本
* 完全支持 C++20 的编译器,例如 clang 或 gcc(Linux 上需要 gcc 11.3.0 或更新版本)
* Visual Studio 或 Visual Studio Build Tool(仅限 Windows)
* 至少 10 GB 的可用磁盘空间
* 初次构建需要 30-60 分钟(后续重建会快得多)

* PyTorch CI 使用 Visual C++ BuildTools,它随 Visual Studio Enterprise、Professional 或 Community 版本提供。你也可以从 https://visualstudio.microsoft.com/visual-cpp-build-tools/ 安装构建工具。默认情况下,Visual Studio Code 不包含这些构建工具。

环境设置示例如下:

  • Linux:
$ source <CONDA_INSTALL_DIR>/bin/activate
$ conda create -y -n <CONDA_NAME>
$ conda activate <CONDA_NAME>
  • Windows:
$ source <CONDA_INSTALL_DIR>\Scripts\activate.bat
$ conda create -y -n <CONDA_NAME>
$ conda activate <CONDA_NAME>
$ call "C:\Program Files\Microsoft Visual Studio\<VERSION>\Community\VC\Auxiliary\Build\vcvarsall.bat" x64

Conda 环境不是必需的。你也可以在标准的虚拟环境中进行 PyTorch 构建,例如使用 uv 等工具创建,前提是你的系统已安装所有无法通过 pip 包获取的必要依赖项(例如 CUDA、MKL)。

NVIDIA CUDA 支持

如果你想编译支持 CUDA 的版本,请从我们的支持矩阵中选择一个受支持的 CUDA 版本,然后安装以下内容:
* NVIDIA CUDA
* NVIDIA cuDNN v8.5 或更高版本
* 与 CUDA 兼容的编译器

注意:你可以参考 cuDNN 支持矩阵 了解 cuDNN 版本与各种受支持的 CUDA、CUDA 驱动程序和 NVIDIA 硬件的对应关系。

如果你想禁用 CUDA 支持,请导出环境变量 USE_CUDA=0
其他可能有用的环境变量可以在 setup.py 中找到。如果 CUDA 安装在非标准位置,请设置 PATH 以便找到你想使用的 nvcc(例如 export PATH=/usr/local/cuda-12.8/bin:$PATH)。

如果你是为 NVIDIA 的 Jetson 平台(Jetson Nano、TX1、TX2、AGX Xavier)构建,Jetson Nano 的 PyTorch 安装说明在此处可用。

AMD ROCm 支持

如果你想编译支持 ROCm 的版本,请安装:
* AMD ROCm 4.0 及以上版本的安装
* ROCm 目前仅支持 Linux 系统。

默认情况下,构建系统期望 ROCm 安装在 /opt/rocm。如果 ROCm 安装在其他目录,必须将 ROCM_PATH 环境变量设置为 ROCm 的安装目录。构建系统会自动检测 AMD GPU 架构。可选地,可以使用 PYTORCH_ROCM_ARCH 环境变量显式设置 AMD GPU 架构

如果你想禁用 ROCm 支持,请导出环境变量 USE_ROCM=0
其他可能有用的环境变量可以在 setup.py 中找到。

Intel GPU 支持

如果你想编译支持 Intel GPU 的版本,请遵循以下步骤:
* Intel GPU 的 PyTorch 先决条件 说明。
* Intel GPU 支持 Linux 和 Windows。

如果你想禁用 Intel GPU 支持,请导出环境变量 USE_XPU=0
其他可能有用的环境变量可以在 setup.py 中找到。

获取 PyTorch 源码

git clone https://github.com/pytorch/pytorch
cd pytorch
# 如果你正在更新现有的代码库
git submodule sync
git submodule update --init --recursive

安装依赖

通用

# 使用上面的“获取 PyTorch 源码”部分克隆源代码后,从 PyTorch 目录运行此命令
pip install --group dev

在 Linux 上

pip install mkl-static mkl-include
# 仅 CUDA:如果需要,为 GPU 添加 LAPACK 支持
# magma 安装:在激活的 conda 环境中运行。指定要安装的 CUDA 版本
.ci/docker/common/install_magma_conda.sh 12.4

# (可选)如果使用 torch.compile 配合 inductor/triton,请安装匹配版本的 triton
# 克隆后从 pytorch 目录运行
# 对于 Intel GPU 支持,请在运行命令前显式设置 `export USE_XPU=1`。
make triton

在 macOS 上

# 仅在 Intel x86 处理器机器上添加此包
pip install mkl-static mkl-include
# 如果需要 torch.distributed,则添加这些包
conda install pkg-config libuv

在 Windows 上

pip install mkl-static mkl-include
# 如果需要 torch.distributed,则添加这些包。
# Windows 上的分布式包支持是原型功能,可能会发生变化。
conda install -c conda-forge libuv=1.51

安装 PyTorch

在 Linux 上

如果你是为 AMD ROCm 编译,请先运行此命令:

# 仅当你为 ROCm 编译时运行此命令
python tools/amd_build/build_amd.py

安装 PyTorch

# conda 环境的 CMake 前缀
export CMAKE_PREFIX_PATH="${CONDA_PREFIX:-'$(dirname $(which conda))/../'}:${CMAKE_PREFIX_PATH}"
python -m pip install --no-build-isolation -v -e .

# 非 conda 环境(例如 Python venv)的 CMake 前缀
# 激活 venv 后调用以下命令
export CMAKE_PREFIX_PATH="${VIRTUAL_ENV}:${CMAKE_PREFIX_PATH}"

在 macOS 上

python -m pip install --no-build-isolation -v -e .

在 Windows 上

如果你想构建遗留的 Python 代码,请参考在遗留代码和 CUDA 上构建

仅 CPU 构建

在此模式下,PyTorch 计算将在你的 CPU 上运行,而不是 GPU 上。

python -m pip install --no-build-isolation -v -e .

关于 OpenMP 的说明:期望的 OpenMP 实现是 Intel OpenMP (iomp)。为了链接到 iomp,你需要手动下载该库并通过调整 CMAKE_INCLUDE_PATHLIB 来设置构建环境。此处的说明是设置 MKL 和 Intel OpenMP 的示例。如果没有这些 CMake 配置,将使用 Microsoft Visual C OpenMP 运行时 (vcomp)。

基于 CUDA 的构建

在此模式下,PyTorch 计算将利用你的 GPU 通过 CUDA 进行更快的数值计算。

构建支持 CUDA 的 PyTorch 需要 NVTX。NVTX 是 CUDA 发行版的一部分,称为 "Nsight Compute"。要将其安装到已安装的 CUDA 上,请再次运行 CUDA 安装程序并勾选相应的复选框。确保在 Visual Studio 之后安装带有 Nsight Compute 的 CUDA。

目前,VS 2017 / 2019 和 Ninja 作为 CMake 的生成器受支持。如果在 PATH 中检测到 ninja.exe,则 Ninja 将用作默认生成器,否则将使用 VS 2017 / 2019。

如果选择 Ninja 作为生成器,最新的 MSVC 将被选作底层工具链。

通常还需要额外的库,如 MagmaoneDNN,又称 MKLDNN 或 DNNLSccache。请参考 [installation-helper](https://github

9 次点击  ∙  0 人收藏  
登录后收藏  
目前尚无回复
0 条回复
About   ·   Help   ·    
OA0 - Omni AI 0 一个探索 AI 的社区
沪ICP备2024103595号-2
Developed with Cursor