目录
1. 环境准备
对比时间,和模型训练的环境相同,可能额外要安装的包是onnxruntime.
pip install onnxruntime # for cpu
pip install onnxruntime-gpu # for gpu
2. 测试过程
直接上代码吧,代码就是最好的解释。
import cv2
import time
import torch
import numpy as np
from torch.nn import DataParallel
from MobileNetV2 import mobilenet_v2
from collections import OrderedDict
from torchvision import transforms as T
import onnxruntime as rt
def torch2onnx(model, save_path):
model.eval()
data = torch.rand(1,3,256,256)
input_names = ['input']
output_names = ['out']
torch.onnx.export(model,
data,
save_path,
export_params=True,
opset_version=11,
input_names=input_names,
output_names=output_names)
print("torch2onnx finish")
def img_process(img_path):
normalize = T.Normalize(mean = [0.5, 0.5, 0.5],
std = [0.5, 0.5, 0.5])
transforms = T.Compose([T.ToTensor(),
normalize])
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (256, 256))
img = transforms(img)
img = img.unsqueeze(0)
return img
def onnx_runtime(img):
sess = rt.InferenceSession("mobilenet-v2.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
t0 = time.time()
for i in range(1000):
pred_onnx = sess.run([output_name], {input_name:np.array(img)})
t1 = time.time()
print("用onnx完成1000次推理消耗的时间:%s" % (t1-t0))
print("用onnx推理的结果如下:")
print(pred_onnx[0].tolist())
def model_load(model_pth):
state_dict = torch.load(model_pth, map_location=device)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k
if name.startswith("module."):
name = name[7:]
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
model.eval()
return model
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = mobilenet_v2().to(device)
model_pth = "./mobilenet-v2.pth"
model = model_load(model_pth)
img = img_process("test.jpg")
t0 = time.time()
for i in range(1000):
outputs = model(img)
t1 = time.time()
print("用pytorch完成1000次推理消耗的时间:%s" % (t1-t0))
print("用pytorch推理的结果如下:")
print(outputs[0].detach().tolist())
print()
torch2onnx(model, "mobilenet-v2.onnx")
onnx_runtime(img)
3. 测试结果与分析
运行上面代码,输出如下(我是用cpu跑的结果):
可以看到用pytorch和onnx的推理结果几乎相同,完全可以接受。然而,用onnx的推理速度是pytorch的好几倍。
OK,就是这么简单~
想要完整代码和模型,请联系博主,下面是微信二维码。