{ "cells": [ { "cell_type": "markdown", "id": "00e6b1b8", "metadata": {}, "source": [ "# 基于Custom原语的自定义算子\n", "\n", "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.9.0/resource/_static/logo_notebook.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.9.0/tutorials/zh_cn/custom_program/operation/mindspore_op_custom_prim.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.9.0/resource/_static/logo_download_code.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.9.0/tutorials/zh_cn/custom_program/operation/mindspore_op_custom_prim.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.9.0/resource/_static/logo_source.svg)](https://atomgit.com/mindspore/docs/blob/r2.9.0/tutorials/source_zh_cn/custom_program/operation/op_custom_prim.ipynb)\n", "\n", "当开发网络遇到内置算子不足以满足需求时，你可以利用MindSpore的Python API中的[Custom](https://www.mindspore.cn/docs/zh-CN/r2.9.0/api_python/ops/mindspore.ops.Custom.html#mindspore-ops-custom)原语方便快捷地进行不同类型自定义算子的定义和使用。\n", "\n", "传统的添加一个自定义算子的方式，需要完成算子原语注册、算子实现、算子信息注册三部分工作。\n", "\n", "其中：\n", "\n", "- 算子原语：定义了算子在网络中的前端接口原型，也是组成网络模型的基础单元，主要包括算子的名称、属性（可选）、输入输出名称、输出shape推理方法、输出数据类型推理方法等信息。\n", "- 算子实现：在Python侧定义函数（JIT类型自定义算子）或C++侧定义类（GPU和CPU自定义算子），描述算子内部计算逻辑的实现。\n", "- 算子信息：描述自定义算子的基本信息，如算子名称、支持的输入输出数据类型、支持的输入输出数据格式和属性等。它是后端做算子选择和映射时的依据。\n", "\n", "相比于传统自定义算子方式，基于`Custom`原语自定义算子具有如下优势：\n", "\n", "- 不同的自定义算子对应的算子原语都是`Custom`原语，无需对每个自定义算子定义一个相应的算子原语。上述提到的三部分工作可以在网络脚本中以统一的接口进行实现，并作为网络表达的一部分，不需要对MindSpore框架进行侵入式修改和重新编译。\n", "- 实现了不同方式自定义算子的接口和使用统一，方便网络开发者根据需要灵活选用不同的自定义方式。\n", "\n", "## 自定义算子分类及适应场景\n", "\n", "基于[Custom](https://www.mindspore.cn/docs/zh-CN/r2.9.0/api_python/ops/mindspore.ops.Custom.html#mindspore-ops-custom)原语的自定义算子支持的算子开发方式包括：pyfunc、aot。不同的算子开发方式适应的场景如下：\n", "\n", "| 算子开发方式 | 开发语言 | 支持平台 | 推荐场景 |\n", "|:-------|:------------------ |:------ |:------------------------|\n", "| [pyfunc](#自定义算子用例) | Python | `CPU` | 快速算法验证的场景 |\n", "| [aot](https://www.mindspore.cn/tutorials/zh-CN/r2.9.0/custom_program/operation/op_custom_aot.html) | Ascend C / CUDA / C++ | `Ascend` `GPU` `CPU` | 需要高性能算子的场景 |\n", "\n", "不同的开发方式使用不同的开发语言实现算子计算逻辑，但是自定义算子的开发流程是一致的，包括算子实现、shape推导、数据类型推理和算子信息注册（可选）。网络开发者可以根据需要选用不同的自定义算子开发方式。在开发者进行自定义算子开发的时候，可以参考如下方式选择对应类型：\n", "\n", "1. 判断后端：如果用户使用的是Ascend和GPU后端，那么就选择aot类型的自定义算子；如果是CPU后端，则根据使用的场景选择；\n", "2. 判断场景：在使用CPU后端的时候，不同的场景对应不同类型的自定义算子推荐：\n", " - 快速验证场景：如果用户希望基于MindSpore做快速验证和开发，对于性能要求不高，或者希望基于Python进行交互，那么选取pyfunc类型的自定义算子；\n", " - 高性能场景：如果用户希望基于MindSpore做高性能计算，或者需要对接第三方算子库，那么选取aot类型自定义算子；\n", "\n", "为了帮助大家更好地使用自定义算子，我们以[pyfunc类型自定义算子](#自定义算子用例)中作为自定义算子的范例展示。此外，我们提供了其他自定义算子的教程包括：\n", "\n", "- aot类型自定义算子：[Ascend平台](https://www.mindspore.cn/tutorials/zh-CN/r2.9.0/custom_program/operation/op_custom_ascendc.html)和[GPU/CPU平台](https://www.mindspore.cn/tutorials/zh-CN/r2.9.0/custom_program/operation/op_custom_aot.html)；\n", "- [自定义算子进阶用法](https://www.mindspore.cn/tutorials/zh-CN/r2.9.0/custom_program/operation/op_custom_adv.html)：算子注册和反向算子。\n", "\n", "> 更多示例可参考MindSpore源码中[tests/st/custom/ops_custom](https://atomgit.com/mindspore/mindspore/tree/v2.9.0/tests/st/custom/ops_custom)下的用例。\n", "\n", "## 自定义算子用例\n", "\n", "为了帮助用户快速入门自定义算子，这里以pyfunc类型自定义算子为例帮助用户理解自定义算子的定义流程。下面基于pyfunc模式定义一个实现sin计算的自定义算子。pyfunc类型的自定义算子使用原生Python语法定义算子实现函数，描述算子内部计算逻辑的实现。网络运行时框架会自动调用此函数。为了表达自定义算子的计算，我们写一个基于numpy的计算正弦函数的Python原生函数。\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "8f447d4c", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "\n", "def sin_by_numpy(x):\n", " return np.sin(x)" ] }, { "cell_type": "markdown", "id": "912957ab", "metadata": {}, "source": [ "然后我们要定义两个函数，一个是张量形状的推导函数（infer_shape），另一个是张量数据类型的推导函数（infer_dtype）。这里要注意：\n", "\n", "- 张量形状的推导函数是输入张量的形状；\n", "- 张量数据类型的推导函数是输入张量的数据类型。" ] }, { "cell_type": "code", "execution_count": 2, "id": "ae19a3fa", "metadata": {}, "outputs": [], "source": [ "def infer_shape(x):\n", "\n", " # 1. 这里的输入x是算子输入张量的形状\n", " # 2. sin函数是逐元素计算，输入的形状和输出的一样\n", " return x\n", "\n", "def infer_dtype(x):\n", "\n", " # 1. 这里的输入x是算子输入张量的数据类型\n", " # 2. sin函数输入的数据类型和输出的一样\n", " return x\n" ] }, { "cell_type": "markdown", "id": "bafa685b", "metadata": {}, "source": [ "下面我们用上面的函数自定义一个算子，其输入包括\n", "\n", "- func：自定义算子的函数表达，这里我们用`sin_by_numpy`函数；\n", "- out_shape：输出形状的推导函数，这里我们用`infer_shape`函数；\n", "- out_dtype：输出数据类型的推导函数，这里我们用`infer_dtype`函数；\n", "- func_type：自定义算子类型，这里我们用`\"pyfunc\"`。" ] }, { "cell_type": "code", "execution_count": 3, "id": "f629190b", "metadata": {}, "outputs": [], "source": [ "from mindspore import ops\n", "\n", "sin_by_numpy_op = ops.Custom(func=sin_by_numpy, # 这里填入自定义算子的函数表达\n", " out_shape=infer_shape, # 这里填入输出形状的推导函数\n", " out_dtype=infer_dtype, # 这里填入输出数据类型的推导函数\n", " func_type=\"pyfunc\" # 这里填入自定义算子类型\n", " )" ] }, { "cell_type": "markdown", "id": "809897ac", "metadata": {}, "source": [ "加上其他环境依赖和算子调用语句，我们获得完整的自定义算子用例如下。" ] }, { "cell_type": "code", "execution_count": 4, "id": "64182df1", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[0. 0.841471 0.19866933 0.29552022 0.38941833]\n" ] } ], "source": [ "import numpy as np\n", "import mindspore as ms\n", "from mindspore import ops\n", "\n", "ms.set_context(mode=ms.GRAPH_MODE)\n", "ms.set_device(device_target=\"CPU\")\n", "\n", "def sin_by_numpy(x):\n", " return np.sin(x)\n", "\n", "def infer_shape(x):\n", " return x\n", "\n", "def infer_dtype(x):\n", " return x\n", "\n", "sin_by_numpy_op = ops.Custom(func=sin_by_numpy,\n", " out_shape=infer_shape,\n", " out_dtype=infer_dtype,\n", " func_type=\"pyfunc\")\n", "input_tensor = ms.Tensor([0, 1, 0.2, 0.3, 0.4], dtype=ms.float32)\n", "result_cus = sin_by_numpy_op(input_tensor)\n", "print(result_cus)" ] }, { "cell_type": "markdown", "id": "d51b50fa", "metadata": {}, "source": [ "我们可以得到结果即为上面输入对应的sin值。" ] }, { "cell_type": "markdown", "id": "91389287", "metadata": {}, "source": [ "如此我们完成一个pyfunc类型自定义算子的定义。对于更多完整的pyfunc类型自定义算子的例子，参见MindSpore源码中的[用例](https://atomgit.com/mindspore/mindspore/blob/v2.9.0/tests/st/custom/ops_custom/test_custom_pyfunc.py)。\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3.8.5 64-bit ('base': conda)", "metadata": { "interpreter": { "hash": "99cb5f0cc54faf3370c23dbaaafe003bb5790f9fd7059c3ab66b2d510ce39296" } }, "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.5-final" } }, "nbformat": 4, "nbformat_minor": 5 }