目录

1. 测试Resnet50 ONNX模型的推理速度

---

1. 测试Resnet50 ONNX模型的推理速度

###############################
# 导出resnet50 模型
# 测试onnx模型推理 cpu 和 GPU 的对比
###############################



import time
import numpy as np
import onnxruntime as ort
import torchvision.models as models
import torch
 
def export_onnx(onnx_path):
    # ResNet50 模型， 保存onnx model 返回 torch model
    resnet50 = models.resnet50(pretrained=False)
    torch.onnx.export(
            resnet50,  # 模型
            torch.randn(1, 3, 224, 224),  # 输入
            onnx_path, # 输出文件名
            opset_version=12, # opset版本 
            input_names=["input"], # 输入变量名
            output_names=["output"], # 输出变量名
            dynamic_axes={"input":{0 : "batch_x"}, "output":{0: "batch_y"}} # 动态轴设置
            )
 
 
class ONNXModel:
    def __init__(self, model_path, use_gpu=False):
        """
        初始化 ONNX 模型推理类。
        
        :param model_path: ONNX 模型文件路径
        :param use_gpu: 是否使用 GPU 推理
        """
        self.model_path = model_path
        self.use_gpu = use_gpu
        self.session = None
        self.load_model()
        self.input_name = self.session.get_inputs()[0].name  # 输入名字
 
    def load_model(self):
        """
        加载 ONNX 模型。
        """
        try:
            print(f"Loading model from {self.model_path}...")
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if self.use_gpu else ['CPUExecutionProvider']
            self.session = ort.InferenceSession(self.model_path, providers=providers)
        except Exception as e:
            print(f"Failed to load model: {e}")
            raise
        
        providers = self.session.get_providers()
        if 'CUDAExecutionProvider' in providers:
            print("✅ ONNX Runtime 正在使用 GPU 推理")
        else:
            print("❌ ONNX Runtime 正在使用 CPU 推理")
 
    def run_inference(self, input_data):
        """
        执行推理并返回结果。
        
        :param input_data: 输入数据，通常为 numpy 数组或字典
        :return: 推理结果
        """
        # 确保输入是字典格式
        input_data = {self.input_name: input_data}
        outputs = self.session.run(None, input_data)
        return outputs[0]
 
 
def eval(onnx_model, n):
    # 推理n 次 返回对比结果
    batch_size = 32
    print("评估中....")
    onnx_input = np.random.randn(batch_size,3,224,224).astype(np.float32)  # ONNX 需要 CPU 数据
    # 1. ONNX 时间
    t1 = time.time()
    for _ in range(n):
        onnx_model.run_inference(onnx_input)
    t2 = time.time()
    t = t2 - t1
    print(f"推理{n}次时间对比:")
    print(f"模型推理时间: {t:.4f} 秒")
    print(f"模型推理平均时间: {t/n:.4f} 秒")
   
 
def test():
    onnx_path = "./resnet50.onnx"
    # export_onnx(onnx_path)


    onnx_model_cpu = ONNXModel(onnx_path, use_gpu=False)
    # 评估推理性能
    eval(onnx_model_cpu, 50)
    
    onnx_model_gpu = ONNXModel(onnx_path, use_gpu=True)
    eval(onnx_model_cpu, 50)

 
# 示例用法
if __name__ == "__main__":
    test()