mindspore.ops.Custom
- class mindspore.ops.Custom(func, bprop=None, out_dtype=None, func_type='hybrid', out_shape=None, reg_info=None)[源代码]
Custom 算子是MindSpore自定义算子的统一接口。用户可以利用该接口自行定义MindSpore内置算子库尚未包含的算子。 根据输入函数的不用,你可以创建多个自定义算子,并且把它们用在神经网络中。 关于自定义算子的详细说明和介绍,包括参数的正确书写,见 教程 。
Warning
这是一个实验性接口,后续可能删除或修改。
Note
不同自定义算子的函数类型(func_type)支持的平台类型不同。每种类型支持的平台如下:
“hybrid”: [“Ascend”, “GPU”, “CPU”].
“akg”: [“Ascend”, “GPU”, “CPU”].
“tbe”: [“Ascend”].
“aot”: [“GPU”, “CPU”].
“pyfunc”: [“CPU”].
“julia”: [“CPU”].
“aicpu”: [“Ascend”].
参数:
func (Union[function, str]) - 自定义算子的函数表达。
out_shape (Union[function, list, tuple]) - 自定义算子的输入的形状或者输出形状的推导函数。默认值:None。
out_dtype (Union[function,
mindspore.dtype
, tuple[mindspore.dtype
]]) - 自定义算子的输入的数据类型或者输出数据类型的推导函数。默认值:None。func_type (str) - 自定义算子的函数类型,必须是[“hybrid”, “akg”, “tbe”, “aot”, “pyfunc”, “julia”, “aicpu”]中之一。默认值:”hybrid”。
bprop (function) - 自定义算子的反向函数。默认值:None。
reg_info (Union[str, dict, list, tuple]) - 自定义算子的算子注册信息。默认值:None。
输入:
input (Union(tuple, list)) - 输入要计算的Tensor。
输出:
Tensor。自定义算子的计算结果。
异常:
TypeError - 如果输入 func 不合法,或者 func 对应的注册信息类型不对。
ValueError - func_type 的值不在列表内。
ValueError - 算子注册信息不合法,包括支持平台不匹配,算子输入和属性与函数不匹配。
- 支持平台:
Ascend
GPU
CPU
样例:
>>> import mindspore.ops as ops >>> import numpy as np >>> from mindspore.ops import CustomRegOp, custom_info_register, DataType, ms_kernel >>> from mindspore.common import dtype as mstype >>> from mindspore.nn import Cell >>> input_x = Tensor(np.ones([16, 16]).astype(np.float32)) >>> input_y = Tensor(np.ones([16, 16]).astype(np.float32)) >>> >>> # Example, func_type = "hybrid" >>> # This is the default func_type in Custom, >>> # and both out_shape and out_dtype can be None(default value). >>> # In this case, the input func must be a function written in the Hybrid DSL >>> # and decorated by @ms_kernel. >>> @ms_kernel ... def outer_product_script(a, b): ... c = output_tensor(a.shape, a.dtype) ... for i0 in range(a.shape[0]): ... for i1 in range(b.shape[1]): ... c[i0, i1] = 0.0 ... for i2 in range(a.shape[1]): ... c[i0, i1] = c[i0, i1] + (a[i0, i2] * b[i2, i1]) ... return c >>> >>> test_op_hybrid = ops.Custom(outer_product_script) >>> output = test_op_hybrid(input_x, input_y) >>> # the result will be a 16 * 16 tensor with all elements 16 >>> print(output.shape) (16, 16) >>> # Example, func_type = "akg" >>> def outer_product(a_1, b_1): ... d = output_tensor(a_1.shape, a_1.dtype) ... for i0 in range(a_1.shape[0]): ... for i1 in range(b_1.shape[1]): ... d[i0, i1] = 0.0 ... for i2 in range(a.shape[1]): ... d[i0, i1] = d[i0, i1] + (a_1[i0, i2] * b_1[i2, i1]) ... return d >>> >>> # Example, func_type = "tbe" >>> square_with_bias_op_info = CustomRegOp() \ ... .fusion_type("OPAQUE") \ ... .attr("bias", "required", "float") \ ... .input(0, "x") \ ... .output(0, "y") \ ... .dtype_format(DataType.F32_Default, DataType.F32_Default) \ ... .dtype_format(DataType.F16_Default, DataType.F16_Default) \ ... .target("Ascend") \ ... .get_op_info() >>> >>> @custom_info_register(square_with_bias_op_info) ... def square_with_bias(input_x, output_y, bias=0.0, kernel_name="square_with_bias"): ... import te.lang.cce ... from te import tvm ... from topi.cce import util ... ... shape = input_x.get("shape") ... dtype = input_x.get("dtype").lower() ... ... shape = util.shape_refine(shape) ... data = tvm.placeholder(shape, name="data", dtype=dtype) ... ... with tvm.target.cce(): ... res0 = te.lang.cce.vmul(data, data) ... res = te.lang.cce.vadds(res0, bias) ... sch = te.lang.cce.auto_schedule(res) ... ... config = {"print_ir": False, ... "name": kernel_name, ... "tensor_list": [data, res]} ... ... te.lang.cce.cce_build_code(sch, config) >>> >>> def test_tbe(): ... square_with_bias = ops.Custom(square_with_bias, out_shape=lambda x, _: x, \ ... out_dtype=lambda x, _: x, func_type="tbe") ... res = self.square_with_bias(input_x, 1.0) ... return res >>> >>> # Example, func_type = "aicpu" >>> resize_bilinear_op_info = CustomRegOp("ResizeBilinear") \ ... .fusion_type("OPAQUE") \ ... .input(0, "input", "required") \ ... .output(1, "output", "required") \ ... .attr("align_corners", "required", "bool") \ ... .attr("cust_aicpu", "optional", "str", "aicpu_kernels") \ ... .dtype_format(DataType.F32_Default, DataType.F32_Default) \ ... .dtype_format(DataType.F16_Default, DataType.F32_Default) \ ... .target("Ascend") \ ... .get_op_info() >>> >>> @custom_info_register(resize_bilinear_op_info) ... def resize_bilinear_aicpu(): ... return >>> >>> def test_aicpu(x): ... resize_bilinear_op = ops.Custom(resize_bilinear_aicpu, out_shape=[1, 1, 9, 9], \ ... out_dtype=mstype.float32, func_type="aicpu") ... res = resize_bilinear_op(x, True, "aicpu_kernels") ... return res >>> >>> # Example, func_type = "aot" >>> def test_aot(x, y, out_shapes, out_types): ... program = ops.Custom("./reorganize.so:CustomReorganize", out_shapes, out_types, "aot") ... out = program(x, y) ... return out >>> >>> # Example, func_type = "pyfunc" >>> def func_multi_output(x1, x2): ... return (x1 + x2), (x1 - x2) >>> >>> test_pyfunc = ops.Custom(func_multi_output, lambda x, _: (x, x), lambda x, _: (x, x), "pyfunc") >>> output = test_pyfunc(input_x, input_y) >>> >>> # Example, func_type = "julia" >>> # julia code: >>> # add.jl >>> # module Add >>> # function add(x, y, z) >>> # z .= x + y >>> # return z >>> # end >>> # end >>> def test_julia(x, y, out_shapes, out_types): ... program = ops.Custom("./add.jl:Add:add", out_shapes, out_types, "julia") ... out = program(x, y) ... return out