npu推理代码

npu 推理代码(这个直接抄的朋友的代码~我试过可以运行,不过环境配置挺费劲的。。)
我一般也用不到npu,这里就放在这儿,以备不时之需吧。

python3 -m bmnetp --model=mnist_jit_0.98.pth --shapes="[(1,1,28,28)]" --net_name="mnist" --target=BM1684 --outdir=.
import sys
import cv2
import numpy as np
import sophon.sail as sail


def usage(name):
    print("使用方法:")
    print("    python3 {} image_path".format(name))

def run(argv):
    # 判断输入参数个数
    if len(argv) == 1:
        usage(argv[0])
        return

    # 读取图像
    image = cv2.imread(argv[1])
    if image is None:
        print("error: 读取图像文件失败")
        return

    # 加载模型
    model = sail.Engine("mnist_0.98.bmodel", 0, sail.IOMode.SYSIO)

    # 获取模型中网络名称
    net_name = model.get_graph_names()[0]

    # 获取网络输入张量名称
    input_name = model.get_input_names(net_name)[0]

    # 图像转灰度
    image_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # 提升维度到 [1, 1, 28, 28]
    image_dims = image_gray[np.newaxis, np.newaxis, :]

    # 组装输入数据
    input_data = {input_name: image_dims}

    # 推理
    output_data = model.process(net_name, input_data)

    # 获取最大分值的类
    output_value = list(output_data.values())[0]
    result = np.argmax(output_value)

    print("识别数字为 {}".format(result))

    return



if __name__ == "__main__":
    run(sys.argv)

你可能感兴趣的:(python深度学习,python)