基于CustomOpBuilder的自定义算子

查看源文件

概述

动态图模式下,网络流程更容易调试,可以支持执行单算子、普通函数和网络,以及单独求梯度等操作。

基于Custom原语的自定义算子虽然可以同时支持静态图和动态图,但是需要定义的内容较多。因此MindSpore针对动态图的自定义算子接入方式做了优化,提供了新的Python API CustomOpBuilder ,在方便用户使用的同时,还能提升动态图自定义算子的执行性能。

用户基于C++接口开发算子,需要定义算子函数体,包括推导并构造输出Tensor,调用执行device算子等功能。定义好算子函数体后,通过pybind11 组件即可将C++函数注册成为Python模块接口。

使用场景示例