ONNX Runtime推理优化实战

# ONNX Runtime量化示例代码
from onnxruntime.quantization import quantize_dynamic, QuantType

# 动态量化FP32模型到INT8
quantize_dynamic(
    model_input='model.onnx',
    model_output='quantized_model.onnx',
    weight_type=QuantType.QUInt8,  # 使用无符号8位整数
    per_channel=False,             # 是否进行通道级量化
    reduce_range=False,            # 是否缩小量化范围
    activation_type=QuantType.QUInt8  # 激活值量化类型
)

print("模型量化完成")

from onnxruntime import OnnxRuntimeError
from onnxruntime.transformers import optimizer
from onnxruntime.transformers.optimizer import optimize_model

# 应用图优化
try:
    optimized_model = optimize_model(
        input='model.onnx',
        optimize_config='default',
        only_onnxruntime=True  # 仅应用ONNX Runtime支持的优化
    )
    optimized_model.save_model_to_file('optimized_model.onnx')
    print("模型图优化完成")
except OnnxRuntimeError as e:
    print(f"图优化失败: {e}")

# 创建虚拟环境
conda create -n onnxrt python=3.9
conda activate onnxrt

# 安装ONNX Runtime
pip install onnxruntime==1.14.0 onnx==1.13.0

# 如需要GPU支持
pip install onnxruntime-gpu==1.14.0

# 安装其他依赖
pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 opencv-python==4.5.5.64

import onnxruntime as ort
import numpy as np
import cv2

def preprocess_image(image_path):
    """图像预处理函数"""
    # 读取图像
    image = cv2.imread(image_path)
    # 调整大小到模型输入尺寸（假设为224x224）
    image = cv2.resize(image, (224, 224))
    # 转换为RGB格式
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 归一化到[0, 1]范围
    image = image.astype(np.float32) / 255.0
    # 添加批次维度
    image = np.expand_dims(image, axis=0)
    # 转换为NCHW格式（如果模型需要）
    image = np.transpose(image, (0, 3, 1, 2))
    return image

def postprocess_output(output):
    """输出后处理函数"""
    # 假设输出是分类概率
    probabilities = np.squeeze(output)
    # 获取最高概率的类别索引
    class_id = np.argmax(probabilities)
    confidence = probabilities[class_id]
    return class_id, confidence

def benchmark_model(session, input_data, num_runs=100):
    """模型性能基准测试"""
    # 预热运行
    for _ in range(10):
        _ = session.run(None, {'input_name': input_data})
    
    # 记录推理时间
    start_time = time.time()
    for _ in range(num_runs):
        _ = session.run(None, {'input_name': input_data})
    end_time = time.time()
    
    avg_inference_time = (end_time - start_time) / num_runs * 1000  # 转换为毫秒
    return avg_inference_time

# 模型推理主流程
if __name__ == "__main__":
    # 初始化ONNX Runtime会话
    # 使用GPU作为执行提供者（如可用）
    providers = ['CUDAExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
    ort_session = ort.InferenceSession('optimized_model.onnx', providers=providers)
    
    # 获取输入和输出节点名称
    input_name = ort_session.get_inputs()[0].name
    output_name = ort_session.get_outputs()[0].name
    
    # 预处理输入图像
    input_image = preprocess_image('test_image.jpg')
    
    # 执行推理
    output = ort_session.run([output_name], {input_name: input_image})
    
    # 后处理输出
    class_id, confidence = postprocess_output(output[0])
    print(f"分类结果: 类别 {class_id}, 置信度 {confidence:.4f}")
    
    # 性能基准测试
    avg_time = benchmark_model(ort_session, input_image, num_runs=100)
    print(f"平均推理时间: {avg_time:.2f} ms")

import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 创建虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出为ONNX格式
torch.onnx.export(
    model,
    dummy_input,
    'resnet18.onnx',
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output']
)

print("模型成功导出为ONNX格式")

import onnx
from onnxruntime import InferenceSession

# 检查ONNX模型结构
onnx_model = onnx.load('resnet18.onnx')
onnx.checker.check_model(onnx_model)
print("ONNX模型验证通过")

# 验证推理正确性
ort_session = InferenceSession('resnet18.onnx')
ort_output = ort_session.run(None, {'input': dummy_input.numpy()})

# 对比PyTorch输出
torch_output = model(dummy_input)
assert np.allclose(ort_output[0], torch_output.detach().numpy(), atol=1e-3)
print("ONNX模型推理结果与PyTorch一致")

from onnxruntime.transformers import optimizer

# 应用优化
optimized_model = optimizer.optimize_model('resnet18.onnx')
optimized_model.save_model_to_file('optimized_resnet18.onnx')

print("模型图优化完成")

from onnxruntime.quantization import quantize_dynamic, QuantType

# 动态量化FP32模型到INT8
quantize_dynamic(
    model_input='optimized_resnet18.onnx',
    model_output='quantized_resnet18.onnx',
    weight_type=QuantType.QUInt8,
    per_channel=True,  # 启用通道级量化
    reduce_range=False
)

print("模型量化完成")