?

目錄

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 torchdef export_onnx(onnx_path):# ResNet50 模型， 保存onnx model 返回 torch modelresnet50 = 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_pathself.use_gpu = use_gpuself.session = Noneself.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}")raiseproviders = 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 = 32print("評估中....")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 - t1print(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()