安装 TorchDynamo

本节描述了如何安装TorchDynamo。 TorchDynamo包含在PyTorch的夜间构建二进制文件中。有关 更多信息，请参见入门指南。

需求

使用 TorchDynamo 必须具备以下前提条件：

• 一个 Linux 或 macOS 环境

• Python 3.8（推荐）。支持并经过测试的 Python 版本为 3.7 到 3.10。请确保本地已安装 Python 的开发版本。

GPU/CUDA 要求

要使用 GPU 后端，特别是 Triton，请确保你本地安装的 CUDA 与正在运行的 PyTorch 版本相匹配。

以下命令安装 GPU PyTorch + TorchDynamo 及其 GPU TorchDynamo 依赖项（适用于 CUDA 11.7）：

pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CPU 要求

CPU TorchDynamo 没有额外的要求。 CPU TorchDynamo 包含在 PyTorch 的夜间构建版本中。 要安装，请运行以下命令：

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

验证安装

如果你是从源代码构建了 PyTorch，那么你可以运行以下命令（从 PyTorch 仓库的根目录） 来检查 TorchDynamo 是否已正确安装：

cd tools/dynamo
python verify_dynamo.py

如果你本地没有 PyTorch 源代码，你也可以从 PyTorch 仓库中复制脚本 (tools/dynamo/verify_dynamo.py) 并在本地运行它。

Docker 安装

我们还在 PyTorch 每日构建的二进制文件中提供了所有必需的依赖项，你可以使用以下命令进行下载：

docker pull ghcr.io/pytorch/pytorch-nightly

对于临时实验，请确保你的容器可以访问所有 GPU：

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash