import onnxruntime as ort
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
# 使用 GPU 执行提供程序配置会话
# 提供程序按顺序尝试 — 先 CUDA,然后 CPU 回退
providers = [
("CUDAExecutionProvider", {
"device_id": 0,
"arena_extend_strategy": "kNextPowerOfTwo",
"gpu_mem_limit": 4 * 1024 * 1024 * 1024, # 4GB 限制
"cudnn_conv_algo_search": "EXHAUSTIVE",
"do_copy_in_default_stream": True,
}),
"CPUExecutionProvider"
]
# 性能相关的会话选项
opts = ort.SessionOptions()
opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
opts.intra_op_num_threads = 8
opts.execution_mode = ort.ExecutionMode.ORT_PARALLEL
# 加载模型
session = ort.InferenceSession(
"resnet50.onnx",
sess_options=opts,
providers=providers
)
print(f"Running on: {session.get_providers()}")
# 准备输入
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
img = Image.open("test_image.jpg").convert("RGB")
img_tensor = transform(img).unsqueeze(0).numpy()
# 运行推理
outputs = session.run(None, {"input": img_tensor})
probabilities = outputs[0][0]
top5_idx = probabilities.argsort()[-5:][::-1]
print("前 5 个预测:", top5_idx, probabilities[top5_idx])