自定义ops.Custom报错TypeError: function output_tensor expects two inputs, but get 1
收藏回复举报
自定义ops.Custom报错TypeError: function output_tensor expects two inputs, but get 1
发表于2023-08-21 15:50:56
0 查看

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend

MindSpore版本: mindspore=1.7.0

执行模式(PyNative/ Graph):不限

Python版本: Python=3.7

操作系统平台: 不限

2 报错信息

2.1 问题描述

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

2.2 报错信息

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

2.3 脚本代码(代码格式,可上传附件)

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)

3 根因分析

******此处由用户补充详细的定位过程******

4 解决方案

cke_7528.png

包含文字方案和最终脚本代码

请将正确的脚本打包并上传附件

我要发帖子