NPU上运行onnxruntime

1 问题背景

在Ascend环境上，使用onnxruntime推理时，报错：

/onnxruntime_src/onnxruntime/core/session/provider_bridge_ort:1193 onnxruntime::Provider& onnxruntime::ProviderLibrary::Get() [ONNXRuntimeError] : 1 : FAIL : Failed to load library libonnxruntime_providers_cuda.so 
with error: libcudart.so.11.0: cannot open shared object file: No such file or directorys

显示的是cuda的组件找不到。由于是Ascend环境，肯定是没有GPU的，很明显是onnxruntime的包装错了。

pip show onnxruntime-gpu

通过如上命令，可以看到环境上确实安装的是gpu版本的。这里先卸载onnxruntime-gpu。

2 解决办法

通过查找onnxruntime社区，可以发现npu已经对onnxruntime进行了适配，参考如下文档：

onnxruntime/docs/execution-providers/community-maintained/CANN-ExecutionProvider.md at gh-pages · microsoft/onnxruntime

那么该如何才能构建onnxruntime适配Ascend上的whl包呢？

这里提供源码构建的办法：

2.1 克隆源码

git clone .git
cd onnxruntime

2.2 构建whl包

注意--use_cann参数

# 先初始化CANN环境
source /usr/local/Ascend/ascend-toolkit/set_env.sh
 
# 执行build脚本
./build.sh --config Release --build_shared_lib --parallel --use_cann --build_wheel

若构建成功，whl包会生成在build/Linux/Release/dist/目录下。安装onnxruntime-cann

pip3 install onnxruntime_cann-x.xx.x-xxxx-xxxxx-xxxxx_xxx_xx.whl

3 推理示例

import numpy as np
import onnxruntime as ort
 
# 模型路径
model = "./resnet34.onnx"
 
# 以CANN作为EP的配置，参数说明可见：.md
providers = [
    ("CANNExecutionProvider", {
        "device_id": 0,                                 
        "arena_extend_strategy": "kNextPowerOfTwo",     
        "npu_mem_limit": 4 * 1024 * 1024 * 1024,
        "enable_cann_graph": True,
    }),
]
 
# options可以用来对推理session进行配置，例如开启profiling功能
options = ort.SessionOptions()
 
# 创建推理session
session = ort.InferenceSession(model, providers=providers, sess_options=options)
 
# 构造纯推理数据
ndata = np.zeros([1, 3, 224, 224], dtype=np.float32)
 
# 执行推理
session.run(["class"], {"image": ndata}])