onnxruntime统计onnx模型gpu运行时间

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)