自定义算子ops.Custom注册问题
收藏回复举报
自定义算子ops.Custom注册问题
t('forum.solved') 已解决
新人帖
发表于2023-06-15 15:14:55
0 查看

MindSpore自定义算子使用ops.Custom注册遇到问题

错误信息:

TypeError: In the function ascend_add written in the Hybrid DSL, function output_tensor expects two inputs, but get 1

运行的代码如下:

import mindspore as ms
from mindspore import nn, ops
from mindspore.ops import ms_hybrid
import numpy as np
import time

@ms_hybrid
def ascend_add(a, b):
    c = output_tensor(a.shape, dtype=a.dtype)
    for arg in grid(a.shape):
        c[arg] = a[arg] + b[arg]
    
    return c

def cpu_add(a, b):
    c = np.zeros(a.shape, dtype=a.dtype)
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            c[i, j] = a[i, j] + b[i, j]
    return c

np_x = np.random.normal(0, 1, [4, 4]).astype(np.float32)
np_y = np.random.normal(0, 1, [4, 4]).astype(np.float32)

start = time.time()
input_x = ms.Tensor(np_x)
input_y = ms.Tensor(np_y)
test_op_akg = ops.Custom(ascend_add)
for _ in range(1000):
    c = test_op_akg(input_x, input_y)
#for _ in range(1000):
end = time.time()
print(end-start)

start = time.time()
for _ in range(1000):
    c = cpu_add(np_x, np_y)
end = time.time()
print(end-start)

开发环境:

mindspore1.7.0-cann5.1.0-py3.7-euler
ascend 910

本帖最后由 匿名用户2024/01/31 15:47:38 编辑

我要发帖子