onnx 模型转换及推理时间对比

目录

1. 环境准备

2. 测试过程

3. 测试结果与分析

---

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，就是这么简单~

想要完整代码和模型，请联系博主，下面是微信二维码。