Ascend Extension for PyTorch插件

简介

本项目开发了名为torch_npu的Ascend Extension for PyTorch插件，使昇腾NPU可以适配PyTorch框架，为使用PyTorch框架的开发者提供昇腾AI处理器的超强算力。

昇腾为基于华为昇腾处理器和软件的行业应用及服务提供全栈AI计算基础设施。您可以通过访问昇腾社区，了解关于昇腾的更多信息。

安装

使用二进制文件进行安装

我们为用户提供可以快速安装torch_npu的whl安装包。在安装torch_npu之前，您需要先安装CANN软件。昇腾辅助软件中有更多关于CANN的版本信息。请参考CANN安装指南获取CANN安装包。

1. 安装PyTorch

通过 pip 安装 PyTorch。

aarch64:

pip3 install torch==2.1.0

x86:

pip3 install torch==2.1.0+cpu  --index-url https://download.pytorch.org/whl/cpu

若使用pip命令安装失败，请使用下载链接或进入PyTorch官方网站进行查询下载对应版本。

2. 安装torch_npu依赖

运行以下命令安装依赖。

pip3 install pyyaml
pip3 install setuptools

3. 安装torch_npu

pip3 install torch-npu==2.1.0.post3

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

使用源代码进行安装

某些特殊场景下，用户可能需要自行编译torch_npu。可以根据昇腾辅助软件表和PyTorch与Python版本配套表选择合适的分支。推荐使用Docker镜像编译torch_npu，可以通过以下步骤获取(建议只挂载工作路径，并避开系统路径，以降低安全风险), 生成的.whl文件路径为./dist/：

1. 克隆torch_npu代码仓

   git clone https://gitee.com/ascend/pytorch.git -b v2.1.0-6.0.rc1 --depth 1

2. 构建镜像

   cd pytorch/ci/docker/{arch} # {arch} for X86 or ARM
   docker build -t manylinux-builder:v1 .

3. 进入Docker容器

   docker run -it -v /{code_path}/pytorch:/home/pytorch manylinux-builder:v1 bash
   # {code_path} is the torch_npu source code path

4. 编译torch_npu

   以Python 3.8 为例。

   cd /home/pytorch
   bash ci/build.sh --python=3.8

卸载

Pytorch框架训练环境的卸载可以参考昇腾官方文档。

torch_npu的卸载只需执行命令：

pip3 uninstall torch_npu

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

入门

前提

运行以下命令初始化CANN环境变量。

# Default path, change it if needed.
source /usr/local/Ascend/ascend-toolkit/set_env.sh

快速验证

可以通过以下样例快速体验昇腾NPU。

import torch
import torch_npu

x = torch.randn(2, 2).npu()
y = torch.randn(2, 2).npu()
z = x.mm(y)

print(z)

PyTorch与Python版本配套表

昇腾辅助软件

PyTorch Extension版本号采用{PyTorch版本}-{昇腾版本}命名规则，前者为PyTorch Extension匹配的PyTorch版本，后者用于匹配CANN版本，详细匹配如下：

建议与交流

欢迎大家为社区做贡献。如果有任何疑问或建议，请提交gitee Issues，我们会尽快回复。感谢您的支持。

分支维护策略

AscendPyTorch版本分支的维护阶段如下：

PyTorch版本维护策略

安全声明

Ascend Extension for PyTorch插件 安全声明

参考文档

有关安装指南、模型迁移和训练/推理教程和API列表等更多详细信息，请参考昇腾社区Ascend Extension for PyTorch。

许可证

Ascend Extension for PyTorch插件使用BSD许可证。详见LICENSE文件。