【MindSpore第七期两日集群营】使用hybrid方式开发自定义算子

继续看视频：https://www.bilibili.com/video/BV1nW4y1b7HP

了解一下hybrid的开发方法：

这里其实需要注意几个问题：

1.需要有ms_kernel装饰器

2.算子内不能直接tensor相加，而是必须用标量相加。

我们先按照例子试试 二维张量的加法函数：

vi self_def2.py

import numpy as np
import mindspore as ms
from mindspore import ops
from mindspore.ops import ms_kernel
from mindspore.nn import Cell

@ms_kernel
def tensor_add_2d(x,y):
    result = output_tensor(x.shape, x.dtype)

    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            result[i,j] = x[i,j] + y[i,j]

    return result

tensor_add_2d_op = ops.Custom(func = tensor_add_2d)

input_tensor_x = ms.Tensor(np.random.normal(0,1,[2,3]).astype(np.float32))
input_tensor_y = ms.Tensor(np.random.normal(0,1,[2,3]).astype(np.float32))
result_cus = tensor_add_2d_op(input_tensor_x,input_tensor_y)
print(result_cus)

执行下看看：

python self_def2.py

参考 https://xihe.mindspore.cn/projects/zichun_ye/cus_ops_repo/blob/cus_ops/cus_ops.py

将其改为三维加法：

import numpy as np
import mindspore as ms
from mindspore import ops
from mindspore.ops import ms_kernel
from mindspore.nn import Cell

@ms_kernel
def tensor_add_3d(x,y):
    result = output_tensor(x.shape, x.dtype)

    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(x.shape[2]):
                result[i,j,k] = x[i,j,k] + y[i,j,k]

    return result

tensor_add_3d_op = ops.Custom(func = tensor_add_3d)

input_tensor_x = ms.Tensor(np.random.normal(0,1,[2,3,4]).astype(np.float32))
print("x\n",input_tensor_x)
input_tensor_y = ms.Tensor(np.random.normal(0,1,[2,3,4]).astype(np.float32))
print("y\n",input_tensor_y)
result_cus = tensor_add_3d_op(input_tensor_x,input_tensor_y)
print("hybrid, tensor_add_3d_op\n",result_cus)复制

执行：

那我们怎么看这个计算结果对不对呢？很简单，按照上一个博客（  使用pyfunc方式开发自定义算子 https://bbs.huaweicloud.com/blogs/375511 ）的方式重做一次，对比结果就可以了：

#pyfunc
ms.set_context(mode=ms.GRAPH_MODE,device_target="CPU")

def infer_shape_py(x,y):
    return x

def infer_dtype_py(x,y):
    return x

tensor_add_3d_py_func = ops.Custom(func = tensor_add_3d,
                out_shape = infer_shape_py,
                out_dtype = infer_dtype_py,
                func_type="pyfunc")

result_pyfunc = tensor_add_3d_py_func(input_tensor_x,input_tensor_y)
print ("pyfunc\n",result_pyfunc)复制

重新执行：

python self_def3.py

比较下 hybrid的结果跟pyfunc的结果好像是一样的。

（全文完，谢谢阅读）