TorchAir 支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。
简介
本项目开发了名为TorchAir(Torch Ascend Intermediate Representation)的扩展库，支持用户基于PyTorch框架和torch_npu插件在昇腾NPU上使用图模式进行推理。 TorchAir继承自PyTorch框架Dynamo模式，将PyTorch的FX图转换为GE计算图，并提供了GE计算图在昇腾NPU的编译与执行的能力。
- 如果您想了解如何使用TorchAir，可以优先通过访问TorchAir图模式使用指南，了解关于TorchAir的更多信息。
- 如果您计划为TorchAir贡献代码，请参考README内容。
安装与卸载
1.安装
- 安装依赖
在安装TorchAir之前，请参考版本配套表和支持型号说明，确保您的硬件能够使用TorchAir，并安装最新昇腾软件栈。
- 编译准备
克隆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模拟执行，可以直接回车跳过。否则，需要输入昇腾处理器开发套件的安装路径（默认安装路径通常为~/Ascend/ascend-toolkit/latest/）。
您可以通过设置
ASCEND_SDK_PATH环境变量指定 SDK 目录 或 设置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>做好权限管控。
快速上手
如果您在配置时未指定Ascend sdk的安装路径，则无法执行CPU上的调试，需要在NPU环境上进行测试。
CPU调试时，需要设置LD_LIBRARY_PATH到生成的fake so文件目录以及sdk目录
tools/env.sh会根据配置生成对应的LD_LIBRARY_PATH（如果您在配置时指定了Ascend sdk安装路径）
source tools/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、torch_npu、CANN以及Python的版本匹配详情如下：
| TorchAir版本 | PyTorch版本 | torch_npu版本 | CANN版本 | Python版本 |
|---|---|---|---|---|
| master（主线） | 2.6.0 + | 在研版本 | 在研版本 | Python3.9.x / Python3.10.x / Python3.11.x / Python3.12.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 |
支持的型号
- Ascend 950PR/Ascend 950DT
- Atlas A3 训练系列产品/Atlas A3 推理系列产品
- Atlas A2 训练系列产品/Atlas A2 推理系列产品
贡献
如果您计划为TorchAir做出贡献，请参考CONTRIBUTING。
联系我们
安全声明
TorchAir安全声明参考SECURITY_README文件。
许可证
TorchAir插件使用BSD许可证。详见LICENSE文件。
免责声明
致TorchAir使用者
- TorchAir提供的模型仅供您用于非商业目的。
- 对于各模型，TorchAir平台仅提示性地向您建议可用于训练的数据集，华为不提供任何数据集，如您使用这些数据集进行训练，请您特别注意应遵守对应数据集的License，如您因使用数据集而产生侵权纠纷，华为不承担任何责任。
- 如您在使用TorchAir模型过程中，发现任何问题（包括但不限于功能问题、合规问题），请在GitCode提交issue，我们将及时审视并解决。
致数据集所有者
如果您不希望您的数据集在TorchAir中的模型被提及，或希望更新TorchAir中的模型关于您的数据集的描述，请在GitCode提交issue，我们将根据您的issue要求删除或更新您的数据集描述。衷心感谢您对TorchAir的理解和贡献。
