pip install onnxruntime-gpu
import torch
import numpy as np
import onnx
import onnxruntime
from onnxruntime.datasets import get_example
import os
import cv2
import time
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
def onnx_out(onnx_path,img):
sess = onnxruntime.InferenceSession(onnx_path,providers=["CPUExecutionProvider"]) #cpu
sess_automation = onnxruntime.InferenceSession(onnx_path) #
sess_gpu = onnxruntime.InferenceSession(onnx_path,providers=["CUDAExecutionProvider"]) #gpu
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
t0 = time.time()
for i in range(100):
pred_onnx = sess.run([output_name], {input_name:np.array(img)})
t1 = time.time()
print("=========>用onnx cpu 完成100次推理消耗的平均时间", round(((t1-t0)/100.0)*1000,3)," ms")
t0 = time.time()
for i in range(100):
pred_onnx = sess_automation.run([output_name], {input_name:np.array(img)})
t1 = time.time()
print("=========>用onnx automation 完成100次推理消耗的平均时间", round(((t1-t0)/100.0)*1000,3)," ms")
torch.cuda.synchronize()
t0 = time.time()
for i in range(100):
pred_onnx = sess_gpu.run([output_name], {input_name:np.array(img)})
torch.cuda.synchronize()
t1 = time.time()
print("=========>用onnx_gpu完成100次推理消耗的平均时间", round(((t1-t0)/100.0)*1000,3)," ms")
def onnx_time():
img = torch.randn(1,3,256,256).detach().cpu().numpy()
# img = cv2.imread("2.jpg")
# mean = torch.Tensor([0.485*255, 0.456*255, 0.406*255]).view(1, -1, 1, 1) #torch.Size([1, 3, 1, 1])
# val = torch.Tensor([0.229*255, 0.224*255, 0.225*255]).view(1, -1, 1, 1) #torch.Size([1, 3, 1, 1])
# img = cv2.resize(img,(256,256)) # wh
# img=np.transpose(img,(2,0,1))
# img = torch.Tensor(img).unsqueeze(0) #nchw
# img = img.sub_(mean).div_(val)
o_out = onnx_out("res50.onnx",img)