十、项目B · 案例四：量化与 ONNX 部署

案例目标：模型量化 → 打包成 ONNX → 编写 load_onnx.py 加载推理。ONNX（Open Neural Network Exchange）是模型界的"通用 PDF"：训练框架各说各话，导出成 ONNX 后任何支持它的运行时（onnxruntime、OpenCV、TensorRT…）都能跑，是工业部署的标准格式。

图 10-1　项目B案例四部署链路：量化 → ONNX 导出 → 运行时加载

10.1 步骤4-1：模型量化（INT8 动态量化最简单可靠）

# quantize_onnx.py —— 先导出再量化（INT8 动态量化对 CNN 检测模型几乎无损）
from ultralytics import YOLO

# ultralytics 一行导出 ONNX（也可视为步骤4-2的一部分）
model = YOLO("runs/sign/exp1/weights/best.pt")
model.export(format="onnx", imgsz=640)     # 得到 best.onnx

# INT8 动态量化
from onnxruntime.quantization import quantize_dynamic, QuantType
quantize_dynamic(
    model_input="best.onnx",
    model_output="best_int8.onnx",
    weight_type=QuantType.QUInt8)          # 权重压成 8bit 整数

10.2 步骤4-3：load_onnx.py 加载与推理（提交的加载代码）

# load_onnx.py —— 纯 onnxruntime 推理：预处理 → session.run → 阈值过滤
import cv2, numpy as np, onnxruntime as ort

# 1. 创建会话；providers 指定优先用 CUDA，没有 GPU 自动退回 CPU
sess = ort.InferenceSession("best_int8.onnx",
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
inp = sess.get_inputs()[0]                       # 拿输入名和期望尺寸
IMG_SIZE = (inp.shape[2], inp.shape[3])          # 如 (640, 640)
CONF_TH, IOU_TH = 0.25, 0.45

def detect(img_path):
    img0 = cv2.imread(img_path)                  # 原图 (BGR)
    # 2. 预处理：letterbox 缩放 → BGR转RGB → HWC转CHW → 归一化0~1 → 加batch维
    r = min(IMG_SIZE[0] / img0.shape[0], IMG_SIZE[1] / img0.shape[1])
    img = cv2.resize(img0, (int(img0.shape[1]*r), int(img0.shape[0]*r)))
    canvas = np.full((*IMG_SIZE, 3), 114, np.uint8)          # 灰边补齐
    canvas[:img.shape[0], :img.shape[1]] = img
    blob = cv2.dnn.blobFromImage(canvas, 1/255.0, swapRB=True).astype(np.float32)

    # 3. 推理：输入名从 sess 里取，别写死
    pred = sess.run(None, {inp.name: blob})[0]  # (1, 84, 8400) yolov8输出
    pred = pred.squeeze().T                      # → (8400, 84)：4框坐标 + 3类分数
    boxes, scores, cls_ids = [], [], []
    for p in pred:
        cls_id = p[4:].argmax(); conf = p[4 + cls_id]
        if conf > CONF_TH:                       # 置信度过滤
            cx, cy, w, h = p[:4]
            # 还原到原图坐标（去 letterbox 缩放）
            boxes.append([int((cx-w/2)/r), int((cy-h/2)/r),
                          int(w/r), int(h/r)])
            scores.append(float(conf)); cls_ids.append(int(cls_id))
    # 4. NMS 去掉重叠框（OpenCV 自带现成实现）
    idx = cv2.dnn.NMSBoxes(boxes, scores, CONF_TH, IOU_TH)
    for i in idx.flatten():
        x, y, w, h = boxes[i]
        cv2.rectangle(img0, (x, y), (x+w, y+h), (0, 200, 0), 2)
        cv2.putText(img0, f"cls{cls_ids[i]} {scores[i]:.2f}", (x, y-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 200, 0), 2)
    return img0, cls_ids

out, labels = detect("test.jpg")
cv2.imwrite("result.jpg", out)
print("检出类别：", labels)

— AI训练师技术交流教程 · 仅供学习交流 —