ONNX Runtime
ONNX Runtime 是一个跨平台 的推理和训练机器学习加速器。为了让应用获得更快的推理速度,我们开发了 SpaceMIT Execution Provider,它会使用扩展 AI 指令进行加速,只需在开发应用时指定 SpaceMITExecutionProvider
。
安装
sudo apt-get update
sudo apt-get install -y onnxruntime python3-spacemit-ort libopencv-dev
验证是否安装成功:
python
>>> import onnxruntime as ort
>>> import spacemit_ort
>>> available_providers = ort.get_available_providers()
>>> print(available_providers)
['SpaceMITExecutionProvider', 'CPUExecutionProvider']
最后输出 ['SpaceMITExecutionProvider', 'CPUExecutionProvider']
表示安装成功。
使用 Python API
安装依赖
sudo apt-get install -y python3-onnx python3-pillow python3-matplotlib python3-opencv
下载模型、标签和图片
我们提供量化好的开源模型,可从 ModelZoo 下载。理论上,[ONNX Model Zoo] 的模型都可以推理,但性能比我们量化的差点。
wget https://archive.spacemit.com/spacemit-ai/ModelZoo/classification.tar.gz
tar -zxf classification.tar.gz
wget -P classification https://archive.spacemit.com/spacemit-ai/ModelZoo/classification/kitten.jpg
wget -P classification https://archive.spacemit.com/spacemit-ai/ModelZoo/classification/synset.txt
导入依赖
确保下面依赖都安装,如果没有错误则继续,警告可忽略。
import onnx
import numpy as np
import onnxruntime as ort
import spacemit_ort
from PIL import Image
import cv2
import matplotlib.pyplot as plt
import time
读图片
def get_image(path, show=False):
with Image.open(path) as img:
img = np.array(img.convert('RGB'))
if show:
plt.imshow(img)
plt.axis('off')
return img
图片预处理
预处理推理图像 -> 缩放到0~1,调整大小到256x256,取224x224的中心裁剪,标准化 图像,转置为NCHW格式,转换为float32并添加一个维度以批处理图像。
def preprocess(img):
img = img / 255.
img = cv2.resize(img, (256, 256))
h, w = img.shape[0], img.shape[1]
y0 = (h - 224) // 2
x0 = (w - 224) // 2
img = img[y0 : y0+224, x0 : x0+224, :]
img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
img = np.transpose(img, axes=[2, 0, 1])
img = img.astype(np.float32)
img = np.expand_dims(img, axis=0)
return img
预测
def predict(path):
img = get_image(path, show=True)
img = preprocess(img)
ort_inputs = {session.get_inputs()[0].name: img}
start = time.time()
preds = session.run(None, ort_inputs)[0]
end = time.time()
preds = np.squeeze(preds)
a = np.argsort(preds)[::-1]
print('time=%.2fms; class=%s; probability=%f' % (round((end-start) * 1000, 2), labels[a[0]], preds[a[0]]))
导入模型、标签并推理
with open('classification/synset.txt', 'r') as f:
labels = [l.rstrip() for l in f]
# Set the number of inference threads
session_options = ort.SessionOptions()
session_options.intra_op_num_threads = 2
# IMPORTANT: Specified SpaceMITExecutionProvider providers
session = ort.InferenceSession('classification/resnet50/resnet50.q.onnx', sess_options=session_options, providers=["SpaceMITExecutionProvider"])
img_path = 'classification/kitten.jpg'
predict(img_path)
结果:
time=95.15ms; class=n02123045 tabby, tabby cat; probability=13.498201
API Overview
更多 API 请查阅官方 API Overview。