简介

本项目开发了名为TorchAir(Torch Ascend Intermediate Representation)的扩展库，支持用户基于PyTorch框架和TorchNPU插件在昇腾NPU上使用图模式进行推理。 TorchAir继承自PyTorch框架Dynamo模式，将PyTorch的FX图转换为GE计算图，并提供了GE计算图在昇腾NPU的编译与执行的能力。

• 如果您想了解如何使用TorchAir，可以优先通过访问TorchAir图模式使用指南，了解关于TorchAir的更多信息。
• 如果您计划为TorchAir贡献代码，请参考README内容。

安装与卸载

1.安装

• 安装依赖

在安装TorchAir之前，请参考版本配套表和支持型号说明，确保您的硬件能够使用TorchAir，并安装最新昇腾软件栈。 建议：执行ut和简单用例请使用docker/Dockerfile中推荐的版本。

• 编译准备

克隆TorchAir代码仓

git clone https://gitcode.com/ascend/torchair.git

下载依赖三方库

cd torchair
git submodule update --init --recursive

配置编译环境，执行配置命令

配置只需要进行一次，用于获取pytorch的编译选项（如当前的torch是否开启ABI）及Ascend sdk的安装路径（如果需要在本地CPU上进行调试）。

cd ./torchair
bash ./configure

默认情况下，执行上述命会弹出如下的交互式会话窗口

您的会话可能有所不同。

Please specify the location of python with available torch 2.1.x installed. [Default is /usr/bin/python3]
(You can make this quiet by set env [TARGET_PYTHON_PATH]):

此时，要求您输入安装了 Torch 2.1 版本的python解释器路径，如果默认路径是正确的，直接回车，否则请输入正确的 python 解释器路径。

您可以通过设置 TARGET_PYTHON_PATH 环境变量，来抑制交互式窗口弹出，但是要确保路径是有效的，否则，仍然会要求您输入正确的 python 解释器路径。

键入后，会耗费几秒钟以确保您的输入是有效的，接着，会弹出下面的交互式窗口

Specify the location of ascend sdk for debug on localhost or leave empty.
(You can make this quiet by set env [ASCEND_SDK_PATH]):

如果您不需要在本地CPU模拟执行，可以直接回车跳过。否则，需要输入昇腾处理器开发套件的安装路径（默认安装路径通常为/usr/local/Ascend/ascend-toolkit/latest/）。

您可以通过设置 ASCEND_SDK_PATH 环境变量指定 Ascend CANN SDK 目录，你后续需要执行UT或者简单用例则需要提前设置,仅编译构建可设置 NO_ASCEND_SDK 环境变量指定不需要 SDK 来抑制交互式窗口弹出。

键入后，等待配置完成。

• 编译

执行以下命令，编译生成TorchAir安装包：

mkdir build
cd build
cmake ..
make torchair -j8

• 安装

编译完成后，会在build/dist/dist/目录下生成名为torchair-{version}-py3-none-any.whl的安装包文件。

您可以直接使用pip安装该安装包，或者使用make命令安装至您configure时指定的python环境中。

make install_torchair

2.卸载

torchair的卸载只需要执行命令：

pip3 uninstall torchair

如需要保存卸载日志，可在pip3 uninstall命令后面加上参数--log <PATH>，并对您指定的目录<PATH>做好权限管控。

快速上手

本的调试(最小CPU调试):

1、在执行./configure设置ASCEND_SDK_PATH如(ASCEND_SDK_PATH=xx/Ascend/ascend-toolkit/latest/)。

2、执行环境设置

source tools/env.sh

npu设备调试：

1、安装torch_npu

torch_npu版本可以根据环境要求安装,请参照[版本配套表]。

pip3 install numpy
pip3 install pyyaml
pip3 install setuptools
pip3 install torch_npu==2.12 # 通过pip安装torch_npu时会自动安装依赖的torch版本

如果安装了真实torch_npu,需要把torch_npu的stubs模块链接环境去除，否则无法正常使用。 如LD_LIBRARY_PATH中去除/workspace/build/stubs设置，如果没有设置请忽略。

2、安装CANN主线版本

npu调试依赖CANN主线版本。如果您是外部用户，可以联系对应的支撑团队获取安装包。

执行前，需要执行 CANN 安装目录下的 set_env.sh 设置 CANN 相关环境变量

例如CANN安装路径为 /usr/local/Ascend/cann，则执行：

source /usr/local/Ascend/cann/set_env.sh

执行以下python脚本快速验证TorchAir基本功能

import torch
import torchair

config = torchair.CompilerConfig()
npu_backend = torchair.get_npu_backend(compiler_config=config)

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        return torch.add(x, y)

model = Model()
model = torch.compile(model, backend=npu_backend, dynamic=False)
x = torch.randn(2, 2)
y = torch.randn(2, 2)
model(x, y)

特性介绍

TorchAir常用特性如下，全量特性介绍请访问docs获取详细教程。

特性功能	功能介绍
日志功能	日志功能
graph dump功能	图dump功能
dynamo export功能	air格式图导出功能
data dump功能	精度数据dump功能
graph fusion功能	用户自定义关闭/开启部分融合算子功能
experimental性能提升功能	试验性质功能，不同功能适用于特定场景，详见参考资料
converter功能拓展	用户自行扩展模型中缺失的converter功能
支持的aten API清单	支持的aten API清单

版本配套表

本项目与PyTorch、TorchNPU、CANN以及Python的版本匹配详情如下：

TorchAir版本	PyTorch版本	TorchNPU版本	CANN版本	Python版本
master（主线）	2.7.1 +	在研版本	在研版本	Python3.10.x
v26.1.0	2.7.1 +	v26.1.0	9.1.0	Python3.10.x / Python3.11.x
v26.0.0	2.7.1 +	v26.0.0	9.0.0	Python3.10.x / Python3.11.x
7.3.0	2.6.0 / 2.7.1 / 2.8.0	7.3.0	8.5.0	Python3.9.x / Python3.10.x / Python3.11.x
7.3.0	2.9.0	7.3.0	8.5.0	Python3.9.x / Python3.10.x / Python3.11.x / Python3.12.x
7.2.0	2.1.0	7.2.0	8.3.RC1	Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.2.0	2.6.0 / 2.7.1 / 2.8.0	7.2.0	8.3.RC1	Python3.9.x / Python3.10.x / Python3.11.x
7.1.0	2.1.0	7.1.0	8.2.RC1	Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.1.0	2.5.1 / 2.6.0	7.1.0	8.2.RC1	Python3.9.x / Python3.10.x / Python3.11.x
7.0.0	2.1.0 / 2.3.1 / 2.4.0	7.0.0	8.1.RC1	Python3.8.x / Python3.9.x / Python3.10.x / Python3.11.x
7.0.0	2.5.1	7.0.0	8.1.RC1	Python3.9.x / Python3.10.x / Python3.11.x

torch_npu的版本选择请参照[版本配套表],下载地址：TorchNPUReleases。

示例: https://gitcode.com/Ascend/pytorch/releases/download/v26.1.0-pytorch2.10.0/torch_npu-2.10.0.post4-cp310-cp310-manylinux_2_28_x86_64.whl

支持的型号

• Ascend 950PR/Ascend 950DT
• Atlas A3 训练系列产品/Atlas A3 推理系列产品
• Atlas A2 训练系列产品/Atlas A2 推理系列产品

贡献

如果您计划为TorchAir做出贡献，请参考CONTRIBUTING。

联系我们

安全声明

TorchAir安全声明参考SECURITY_README文件。

许可证

TorchAir插件使用BSD许可证。详见LICENSE文件。

---

免责声明

致TorchAir使用者

1. TorchAir提供的模型仅供您用于非商业目的。
2. 对于各模型，TorchAir平台仅提示性地向您建议可用于训练的数据集，华为不提供任何数据集，如您使用这些数据集进行训练，请您特别注意应遵守对应数据集的License，如您因使用数据集而产生侵权纠纷，华为不承担任何责任。
3. 如您在使用TorchAir模型过程中，发现任何问题（包括但不限于功能问题、合规问题），请在GitCode提交issue，我们将及时审视并解决。

致数据集所有者

如果您不希望您的数据集在TorchAir中的模型被提及，或希望更新TorchAir中的模型关于您的数据集的描述，请在GitCode提交issue，我们将根据您的issue要求删除或更新您的数据集描述。衷心感谢您对TorchAir的理解和贡献。