--- title: PyTorch框架 description: "通过PyTorch框架进行模型的训练、推理时，会调用很多算子进行计算。开发者开发的自定义算子如果需要集成部署到PyTorch框架，有如下几种方式：" url: https://www.hiascend.com/document/detail/zh/canncommercial/latest/programug/Ascendcopdevg/atlas_ascendc_10_0057.html sourcePath: /source/zh/canncommercial/900/programug/Ascendcopdevg/atlas_ascendc_10_0057.html indexId: b3dc45441cfdb28b96ab9509421a5d5df3277634de79abdee28bbf70d06c92a880 --- # PyTorch框架 通过PyTorch框架进行模型的训练、推理时，会调用很多算子进行计算。开发者开发的自定义算子如果需要集成部署到PyTorch框架，有如下几种方式： - Kernel直调：通过适配torch.library或Pybind注册自定义算子，可以实现PyTorch框架调用算子Kernel程序。 - 单算子API调用：该模式下的适配插件开发流程和具体样例请参见《Ascend Extension for PyTorch 框架特性指南(https://www.hiascend.com/document/detail/zh/Pytorch/2600/ptmoddevg/Frameworkfeatures/docs/zh/framework_feature_guide_pytorch/overview.md)》中的“基于OpPlugin算子适配开发(https://www.hiascend.com/document/detail/zh/Pytorch/2600/ptmoddevg/Frameworkfeatures/docs/zh/framework_feature_guide_pytorch/adaptation_overview_opplugin.md)”章节。 - 图模式调用：自定义算子在Pytorch图模式下的适配开发指导请参见《PyTorch图模式使用指南(TorchAir)(https://www.hiascend.com/document/detail/zh/Pytorch/2600/modthirdparty/torchairuseguide/docs/zh/overview.md)》中的“自定义算子入图(https://www.hiascend.com/document/detail/zh/Pytorch/2600/modthirdparty/torchairuseguide/docs/zh/custom_op_graph/overview.md)”章节。 图1 Pytorch框架部署方式 本节主要提供通过torch.library与Pybind注册自定义算子并实现PyTorch框架调用算子Kernel程序的指导。 - torch.library是用于扩展PyTorch核心算子库的API集合。它允许开发者创建新的算子、并为其提供自定义实现。 - Pybind是一个开源的C++和Python之间的桥接工具，它旨在使C++代码能够无缝地集成到Python环境中。 Pybind适用于快速将C++函数暴露给Python，实现高效接口绑定。但其生成的算子无法被PyTorch的算子系统识别，不具备schema定义与图追踪能力，因此不支持torch.compile优化。相比之下，torch.library提供了与PyTorch核心算子系统深度集成的机制，支持算子注册、schema定义和图追踪能力，是支持torch.compile的必要条件。开发者可根据需求选择对应方式。 #### torch.library 下面代码以add_custom（Add自定义算子为例）算子为例，介绍通过torch.library如何调用算子Kernel程序，文档中仅介绍核心步骤，完整样例请参考torch.library样例(https://gitcode.com/cann/asc-devkit/tree/9.0.0/examples/01_simd_cpp_api/02_features/02_framework/00_pytorch/torch_library)。 1. 环境准备。 除了按照环境准备 进行CANN软件包的安装，还需要安装以下依赖： - “安装PyTorch(https://www.hiascend.com/document/detail/zh/Pytorch/2600/configandinstg/instg/docs/zh/installation_guide/installation_via_binary_package.md)” 2. 实现NPU上的自定义算子。 包括算子Kernel实现，并使用<<<>>>接口调用算子核函数完成指定的运算。样例中的c10_npu::getCurrentNPUStream接口用于获取当前npu流，返回值类型NPUStream，使用方式请参考《Ascend Extension for PyTorch 自定义API参考(https://www.hiascend.com/document/detail/zh/Pytorch/2600/apiref/torchnpuCustomsapi/docs/zh/custom_APIs/overview.md)》中的“（beta）c10_npu::getCurrentNPUStream(https://www.hiascend.com/document/detail/zh/Pytorch/2600/apiref/torchnpuCustomsapi/docs/zh/custom_APIs/cpp/（beta）c10_npu-getCurrentNPUStream.md)”章节。 需要注意的是，本样例的输入x，y的内存是在外层的Python调用脚本中分配的。 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 namespace ascendc_ops { at::Tensor ascendc_add(const at::Tensor& x, const at::Tensor& y) { // 运行资源申请，通过c10_npu::getCurrentNPUStream()的函数获取当前NPU上的流 auto aclStream = c10_npu::getCurrentNPUStream().stream(false); // 分配Device侧输出内存 at::Tensor z = at::empty_like(x); uint32_t numBlocks = 8; uint32_t totalLength = 1; for (uint32_t size : x.sizes()) { totalLength *= size; } // 用<<<>>>接口调用核函数完成指定的运算 add_custom<<>>((uint8_t*)(x.mutable_data_ptr()), (uint8_t*)(y.mutable_data_ptr()), (uint8_t*)(z.mutable_data_ptr()), totalLength); // 将Device上的运算结果拷贝回Host并释放申请的资源 return z; } } // namespace ascendc_ops 3. 自定义算子的注册。 PyTorch提供TORCH_LIBRARY宏作为自定义算子注册的核心接口，用于创建并初始化自定义算子库，注册后在Python侧可以通过torch.ops.namespace.op_name方式进行调用。TORCH_LIBRARY_IMPL用于将算子逻辑绑定到特定的DispatchKey（PyTorch设备调度标识），针对NPU设备，需要将算子实现注册到PrivateUse1这一专属的DispatchKey上。 1 2 3 4 5 6 7 8 9 10 11 // 注册算子到torch.library TORCH_LIBRARY(ascendc_ops, m) { m.def("ascendc_add(Tensor x, Tensor y) -> Tensor"); } // 注册PrivateUse1实现，NPU设备 TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m) { m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add)); } 4. 编译生成算子动态库。 5. 使用Python测试脚本进行测试。 在add_custom_test.py中，首先通过torch.ops.load_library加载生成的自定义算子库，调用注册的ascendc_add函数，并通过对比NPU输出与CPU标准加法结果来验证自定义算子的数值正确性。 #### Pybind 下面代码以add_custom算子为例，介绍通过Pybind方式实现Pytorch脚本中调用自定义算子的流程。文档中仅介绍核心步骤，完整样例请参考Pybind样例(https://gitcode.com/cann/asc-devkit/tree/9.0.0/examples/01_simd_cpp_api/02_features/02_framework/00_pytorch/pybind)。 1. 环境准备。 除了按照环境准备 进行CANN软件包的安装，还需要安装以下依赖： - “安装PyTorch(https://www.hiascend.com/document/detail/zh/Pytorch/2600/configandinstg/instg/docs/zh/installation_guide/installation_via_binary_package.md)” - 安装pybind11 ``` pip3 install pybind11 expecttest ``` 2. 实现NPU上的自定义算子。 包括算子Kernel实现，并使用<<<>>>接口调用算子核函数完成指定的运算。样例中的c10_npu::getCurrentNPUStream接口用于获取当前npu流，返回值类型NPUStream，使用方式请参考《Ascend Extension for PyTorch 自定义API参考(https://www.hiascend.com/document/detail/zh/Pytorch/2600/apiref/torchnpuCustomsapi/docs/zh/custom_APIs/overview.md)》中的“（beta）c10_npu::getCurrentNPUStream(https://www.hiascend.com/document/detail/zh/Pytorch/2600/apiref/torchnpuCustomsapi/docs/zh/custom_APIs/cpp/（beta）c10_npu-getCurrentNPUStream.md)”章节。 需要注意的是，本样例的输入x，y的内存是在Python调用脚本中分配的。 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 // Pybind和PyTorch调用所需的头文件 #include #include #include "torch_npu/csrc/core/npu/NPUStream.h" // Kernel侧实现需要的头文件 #include "kernel_operator.h" ... namespace ascendc_ops { at::Tensor ascendc_add(const at::Tensor& x, const at::Tensor& y) { // 运行资源申请，通过c10_npu::getCurrentNPUStream()的函数获取当前NPU上的流 auto aclStream = c10_npu::getCurrentNPUStream().stream(false); // 分配Device侧输出内存 at::Tensor z = at::empty_like(x); uint32_t numBlocks = 8; uint32_t totalLength = 1; for (uint32_t size : x.sizes()) { totalLength *= size; } // 用<<<>>>接口调用核函数完成指定的运算 add_custom<<>>((uint8_t*)(x.mutable_data_ptr()), (uint8_t*)(y.mutable_data_ptr()), (uint8_t*)(z.mutable_data_ptr()), totalLength); // 将Device上的运算结果拷贝回Host并释放申请的资源 return z; } } // namespace ascendc_ops 3. 定义Pybind模块，将C++函数封装成Python函数。PYBIND11_MODULE是Pybind11库中的一个宏，用于定义一个Python模块。它接受两个参数，第一个参数是封装后的模块名，第二个参数是一个Pybind11模块对象，用于定义模块中的函数、类、常量等。通过调用m.def()方法，可以将上一步骤中函数ascendc_ops::ascendc_add转成Python函数ascendc_add，使其可以在Python代码中被调用。 1 2 3 4 5 PYBIND11_MODULE(ascendc_ops, m)// 模块名ascendc_ops，模块对象m { m.doc() = "add_custom pybind11 interfaces";// optional module docstring m.def("ascendc_add", &ascendc_ops::ascendc_add, "");// 将函数ascendc_add与Pybind模块进行绑定 } 4. 编译生成算子动态库。 5. 在Python调用脚本中，使用torch接口生成随机输入数据并分配内存，通过导入封装的自定义模块ascendc_ops，调用自定义模块ascendc_ops中的run_add_custom函数，从而在NPU上执行算子。