alexnet pytorch模型和onnx模型速度对比

qq_652530495

已于 2022-09-22 08:38:21 修改

阅读量1k

点赞数 1

分类专栏： pytorch onnx 文章标签： pytorch python 深度学习

于 2022-09-21 17:20:02 首次发布

本文链接：https://blog.csdn.net/qq_22526061/article/details/126976847

版权

pytorch 同时被 2 个专栏收录

4 篇文章 0 订阅

订阅专栏

onnx

1 篇文章 0 订阅

订阅专栏

# -*- coding: utf-8 -*-
from tqdm import tqdm
import pandas as pd
import time
import torch
import torchvision
import numpy as np
import onnxruntime as ort


def save_onnx(path, device):
    dummy_input = torch.randn(3, 3, 224, 224).to(device)
    model = torchvision.models.alexnet(pretrained=False).to(device)
    input_names = ["actual_input_1"] + ["learned_%d" % i for i in range(16)]
    output_names = ["output1"]
    torch.onnx.export(model, dummy_input, path, verbose=False, input_names=input_names, output_names=output_names,
                      dynamic_axes={'actual_input_1': [0], 'output1': [0]})
    # print(model)
    return model


def torch_t(model, batch_size, device):
    s = time.time()
    input_ = torch.randn(batch_size, 3, 224, 224).to(device)
    for i in range(epoch):
        model(input_)
    cost = time.time() - s
    return round(cost, 2)


def onnx_t_fun(path, batch_size, device):
    if device == 'cpu':
        ort_session = ort.InferenceSession(path, providers=['CPUExecutionProvider'])
    else:
        ort_session = ort.InferenceSession(path, providers=['CUDAExecutionProvider'])
    s = time.time()
    input_ = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)
    for i in range(epoch):
        ort_session.run(
            None,
            {"actual_input_1": input_},
        )
    cost = time.time() - s
    return round(cost, 2)


if __name__ == '__main__':

    batch_size_list = [2 ** i for i in range(8)]
    device_list = ['cpu', torch.device(0)]
    tuples = [(device, batch_size) for device in device_list for batch_size in batch_size_list]
    index = pd.MultiIndex.from_tuples(tuples)
    epoch = 100
    path = "alexnet.onnx"
    df = pd.DataFrame(columns=['torch', 'onnx'], index=index)
    for batch_size in tqdm(batch_size_list):
        for device in device_list:
            model = save_onnx(path, device)
            cost_torch = torch_t(model, batch_size, device)
            cost_onnx = onnx_t_fun(path, batch_size, device)
            df.loc[(device, batch_size), :] = [cost_torch, cost_onnx]
    print(df)

torch onnx
cpu 1 1.2 0.89
2 2.47 1.15
4 2.86 1.68
8 3.67 2.7
16 5.93 4.72
32 9.47 8.85
64 17.33 17.26
128 32.77 34.27
cuda:0 1 0.07 0.42
2 0.08 0.65
4 0.12 0.6
8 0.2 0.27
16 0.24 0.4
32 0.38 0.89
64 0.69 1.72
128 1.26 3.12