Skip to main content

4.2 Python 推理示例

本文以图像分类为例，介绍如何在 Python 环境中完成模型推理。图像分类是计算机视觉中的基础任务，旨在识别图像中的主要对象类别（如猫、狗、汽车等）。本示例基于 MobileNetV2 模型，并通过 onnxruntime 与 spacemit-ort 加速库在 SpaceMIT 平台上实现高性能推理。

克隆代码​

git clone https://gitee.com/bianbu/spacemit-demo.git ~

安装依赖​

建议使用 Python 虚拟环境：

python -m venv .venv
source .venv/bin/activate

使用 SpaceMIT 私有源安装必要依赖：

pip install opencv-python --index-url https://git.spacemit.com/api/v4/projects/33/packages/pypi/simple
pip install onnxruntime --index-url https://git.spacemit.com/api/v4/projects/33/packages/pypi/simple
pip install spacemit-ort --index-url https://git.spacemit.com/api/v4/projects/33/packages/pypi/simple

⚠️ 注意：spacemit-ort 依赖必须在 onnxruntime 导入之后使用，否则可能导致推理设备识别异常。

运行测试​

进入项目示例目录，执行推理脚本：

cd ~/spacemit_demo/examples/CV/mobilenet_v2/python
python test_mobilenet_v2.py

示例输出：

Final result: class=n02123045 tabby, tabby cat;

推理流程与代码解析​

导入加速库​

import onnxruntime
import spacemit_ort  # 必须在 onnxruntime 之后导入

图像加载与预处理​

from PIL import Image
import numpy as np
import cv2

def preprocess(img):
    img = img / 255.0  # 归一化至 [0, 1]
    img = cv2.resize(img, (256, 256))  # 调整大小
    y0, x0 = (256 - 224) // 2, (256 - 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, (2, 0, 1))  # HWC → CHW
    img = np.expand_dims(img.astype(np.float32), axis=0)  # 添加 batch 维度
    return img

# 加载并预处理图像
img = Image.open(args.image_path).convert('RGB')
img = preprocess(np.array(img))

模型加载与推理 Session 设置​

session_options = onnxruntime.SessionOptions() # onnxruntime session创建
session_options.intra_op_num_threads = 4  # 设置线程数（建议为 4）

# 加载模型，设置推理session
session = onnxruntime.InferenceSession(
    args.model,
    sess_options=session_options,
    providers=["SpaceMITExecutionProvider"]
)

模型推理执行​

input_name = session.get_inputs()[0].name # 获取输入节点的name
output_name = session.get_outputs()[0].name # 获取输出节点的name

# 执行推理，结果保存在result里
result = session.run([output_name], {input_name: img})[0]

推理结果后处理​

# 获取 top-5 分类索引，按概率从高到低排序
top_k = result.argsort()[-5:][::-1]