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)