推理部署 ONNX Runtime
过程：训练---> 转换 ---> onnx模型 ---> 部署推理
1.先看一个简单案例
在这之前，你需要安装包 pytorch、numpy、onnxruntime 或者 onnxruntime-gpu，你可以下面命令安装，pytorch请查询百度
pip install numpy onnxruntime
1.1 导出onnx模型
案例实现两个数相加的简单模型，并且导出onnx模型。
import numpy as np
import torch
import onnxruntime
# 定义模型
def model():
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, x, y):
return x.add(y)
return Model()
# 创建模型
def create_model(type: torch.dtype = torch.float32):
sample_x = torch.ones(3, dtype=type) # 创建一个张量
sample_y = torch.zeros(3, dtype=type) # 创建一个张量
# 导出onnx格式的模型
torch.onnx.export(model(),
(sample_x, sample_y), # 两个输入
r'./model.onnx', # 输出文件名
opset_version=12, # opset版本
input_names=["x", "y"], # 输入名称
output_names=["z"]) # 输出名称
# torch.onnx.export(model(),
# (sample_x, sample_y),
# r'./model.onnx',
# input_names=["x", "y"],
# output_names=["z"],
# dynamic_axes={"x":{0 : "array_length_x"}, "y":{0: "array_length_y"}}) # 动态轴设置
if __name__ == '__main__':
create_model()
1.2 查看onnx 模型
可以在这个网站上查看，选择打开onnx文件 ！Netron
如下图所示！反映了模型的输入和输出的名字，shape，类型等信息
1.3 python 实现推理
import numpy as np
import onnxruntime as ort
# 创建onnx session
def create_session(path):
providers = ['CPUExecutionProvider'] # CPU
# providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] # GPU
sess = ort.InferenceSession(path, providers=providers)
get_onnx_info(sess) # 打印模型信息
return sess
# 推理
def run(x: np.array, y: np.array) -> np.array:
z = session.run(["z"], {"x": x, "y": y})
return z
def get_onnx_info(sess):
input_tensors = sess.get_inputs()
for input_tensor in input_tensors: # 因为可能有多个输入，所以为列表
input_info = {
"name" : input_tensor.name,
"type" : input_tensor.type,
"shape": input_tensor.shape,
}
print("input_info:")
print(input_info)
out_tensors = sess.get_outputs()
for out_tensor in out_tensors: # 因为可能有多个输入，所以为列表
out_info = {
"name" : out_tensor.name,
"type" : out_tensor.type,
"shape": out_tensor.shape,
}
print("out_info:")
print(out_info)
if __name__ == "__main__":
onnx_path = "./model.onnx"
session = create_session(onnx_path) # 创建session
y = run(x=np.float32([1.0, 2.0, 3.0]),y=np.float32([4.0, 5.0, 6.0])) # 推理 两个输入，实现相加
print(y)2. 评估onnx 推理
这个案例中，采用随机模型参数，没有进行训练。在cpu上测试
2.1 导出resnet 模型
import torchvision.models as models
import torch
# 加载预训练的 ResNet50 模型
resnet50 = models.resnet50(pretrained=False)
torch.onnx.export(
resnet50, # 模型
torch.randn(1, 3, 224, 224), # 输入
"./resnet50.onnx", # 输出文件名
opset_version=12, # opset版本
input_names=["input"], # 输入变量名
output_names=["output"] # 输出变量名
)
2.2 模型对比
- 误差评估
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"] # 输出变量名
)
return resnet50
class ONNXModel:
def __init__(self, model_path):
"""
初始化 ONNX 模型推理类。
:param model_path: ONNX 模型文件路径
"""
self.model_path = model_path
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}...")
start_time = time.time()
self.session = ort.InferenceSession(self.model_path)
end_time = time.time()
print(f"Model loaded successfully in {end_time - start_time:.4f} seconds.")
except Exception as e:
print(f"Failed to load model: {e}")
raise
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_difference(torch_model,onnx_mode, n):
# 推理n 次 返回对比结果
print("评估中....")
res_list = []
for _ in range(n):
torch_input = torch.randn(1, 3, 224, 224)
onnx_input= np.asarray(torch_input).astype(np.float32)
onnx_output = onnx_mode.run_inference(onnx_input)
torch_output = torch_model(torch_input).detach().numpy()
res_list.append(np.abs(onnx_output - onnx_output).max())
print(f"平均误差为：{np.mean(np.array(res_list))}")
def test():
onnx_path = "./resnet50.onnx"
torch_model = export_onnx(onnx_path) # 替换为你的 ONNX 模型路径
onnx_model = ONNXModel(onnx_path)
eval_difference(torch_model,onnx_model, 10)
# 示例用法
if __name__ == "__main__":
test()- 时间评估
cpu上
推理50次时间对比:
ONNX 模型推理时间: 0.7116 秒
Torch 模型推理时间: 2.9859 秒
ONNX 模型推理平均时间: 0.0142 秒
Torch 模型推理平均时间: 0.0597 秒
Torch 模型推理时间/ONNX 模型推理时间: 4.1960 秒
ONNX 模型推理时间/Torch 模型推理时间: 0.2383 秒
Torch 模型推理时间 - ONNX 模型推理时间: 2.2743 秒
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"] # 输出变量名
)
return resnet50
class ONNXModel:
def __init__(self, model_path):
"""
初始化 ONNX 模型推理类。
:param model_path: ONNX 模型文件路径
"""
self.model_path = model_path
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}...")
start_time = time.time()
self.session = ort.InferenceSession(self.model_path)
end_time = time.time()
print(f"Model loaded successfully in {end_time - start_time:.4f} seconds.")
except Exception as e:
print(f"Failed to load model: {e}")
raise
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_difference(torch_model,onnx_mode, n):
# 推理n 次 返回对比结果
print("评估中....")
torch_input = torch.randn(1, 3, 224, 224)
onnx_input= np.asarray(torch_input).astype(np.float32)
# 1. onnx 时间
t1 = time.time()
for _ in range(n):
onnx_mode.run_inference(onnx_input)
t2 = time.time()
t_onnx = t2 - t1
# 2. torch 时间
t1 = time.time()
for _ in range(n):
torch_model(torch_input).detach().numpy()
t2 = time.time()
t_torch = t2 - t1
print(f"推理{n}次时间对比:")
print(f"ONNX 模型推理时间: {t_onnx:.4f} 秒")
print(f"Torch 模型推理时间: {t_torch:.4f} 秒")
print(f"ONNX 模型推理平均时间: {t_onnx/n:.4f} 秒")
print(f"Torch 模型推理平均时间: {t_torch/n:.4f} 秒")
print(f"Torch 模型推理时间/ONNX 模型推理时间: {t_torch/t_onnx:.4f} 秒")
print(f"ONNX 模型推理时间/Torch 模型推理时间: {t_onnx/t_torch:.4f} 秒")
print(f"Torch 模型推理时间 - ONNX 模型推理时间: {t_torch - t_onnx:.4f} 秒")
def test():
onnx_path = "./resnet50.onnx"
torch_model = export_onnx(onnx_path) # 替换为你的 ONNX 模型路径
onnx_model = ONNXModel(onnx_path)
eval_difference(torch_model,onnx_model, 50)
# 示例用法
if __name__ == "__main__":
test()2.3 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()
浙公网安备 33010602011771号