pytorch.onnx.export方法参数详解,以及onnxruntime-gpu推理性能测试

Torch.onnx.export执行流程:

1、如果输入到torch.onnx.export的模型是nn.Module类型,则默认会将模型使用torch.jit.trace转换为ScriptModule

2、使用args参数和torch.jit.trace将模型转换为ScriptModule,torch.jit.trace不能处理模型中的循环和if语句

3、如果模型中存在循环或者if语句,在执行torch.onnx.export之前先使用torch.jit.script将nn.Module转换为ScriptModule

4、模型转换成onnx之后,预测结果与之前会有稍微的差别,这些差别往往不会改变模型的预测结果,比如预测的概率在小数点之后五六位有差别。

Onnx模型导出,并能够处理动态的batch_size:

Torch.onnx.export导出模型:

检查导出的模型:

onnxruntime执行导出的onnx模型:

onnxruntime-gpu推理性能测试:

备注:安装onnxruntime-gpu版本时,要与CUDA以及cudnn版本匹配

网络结构:修改Resnet18输入层和输出层,输入层接收[N, 1, 64, 1001]大小的数据,输出256维

测试数据(重复执行10000次,去掉前两次的模型warmup):

输入数据batch_size = 1:[1, 1, 64, 1001]
        pytorch:            2.59 ms
        onnxruntime-gpu:    2.28 ms
        性能提升:           12%
        GPU峰值使用率:       95% vs 50% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

输入数据batch_size = 2:[2, 1, 64, 1001]
        pytorch:            2.92 ms
        onnxruntime-gpu:    2.73 ms
        性能提升:           6.5%
        GPU峰值使用率:       100% vs 41% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

输入数据batch_size = 4:[4, 1, 64, 1001]
        pytorch:            3.93 ms
        onnxruntime-gpu:    3.94 ms
        性能提升:           0%
        GPU峰值使用率:       100% vs 33% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

输入数据batch_size = 8:[8, 1, 64, 1001]
        pytorch:            6.84 ms
        onnxruntime-gpu:    21 ms
        性能提升:           -207%
        GPU峰值使用率:       100% vs 61% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

输入数据batch_size = 16:[16, 1, 64, 1001]
        pytorch:            13.85 ms
        onnxruntime-gpu:    11.41 ms
        性能提升:           17.6%
        GPU峰值使用率:       100% vs 29% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

输入数据batch_size = 32:[32, 1, 64, 1001]
        pytorch:            22.64 ms
        onnxruntime-gpu:    23.56 ms
        性能提升:           0%
        GPU峰值使用率:       100% vs 28% (Tesla P40, pytorch在前)
        CPU使用率:           单核100%

结论:onnxruntime可以提升模型推理速度,但是不擅长处理批量数据,不知道能不能这样理解?还是说只对于我这个网络模型是这个情况,没找到规律,也没找到参考文档,暂时还没搞清楚。关于这个问题,github上找到了这几个问题可以参考:

Gap in inference time between onnxruntime and torch vanishes when increasing the batch size · Issue #9660 · microsoft/onnxruntime · GitHub

https://github.com/microsoft/onnxruntime/issues/2796

https://towardsdatascience.com/an-empirical-approach-to-speedup-your-bert-inference-with-onnx-torchscript-91da336b3a41

 

onnx模型导出及onnxruntime推理完整代码:

import torch
import torchvision
import onnx
import onnxruntime
import numpy as np
import os

# 设置pytorch下载的预训练模型保存位置
os.environ["TORCH_HOME"] = "./pretrained_models"


def pytorch_2_onnx():
    """
    将pytorch模型导出为onnx,导出时pytorch内部使用的是trace或者script先执行一次模型推理,然后记录下网络图结构
    所以,要导出的模型要能够被trace或者script进行转换
    :return:
    """
    # 加载预训练模型
    model = torchvision.models.alexnet(pretrained=True)
    print(model)

    model_path = "alexnet.onnx"

    # pytorch转换为onnx内部使用trace或者script,需要提供一组输入数据执行一次模型推理过程,然后进行trace记录
    dummy_input = torch.randn(4, 3, 224, 224, device="cpu")
    input_names = ["input_data"] + ["learned_%d" % i for i in range(16)]
    output_names = ["output_data"]

    torch.onnx.export(
        model,  # pytorch网络模型
        dummy_input,    # 随机的模拟输入
        model_path,     # 导出的onnx文件位置
        export_params=True,     # 导出训练好的模型参数
        verbose=10,             # debug message
        training=torch.onnx.TrainingMode.EVAL,  # 导出模型调整到推理状态,将dropout,BatchNorm等涉及的超参数固定
        input_names=input_names,    # 为静态网络图中的输入节点设置别名,在进行onnx推理时,将input_names字段与输入数据绑定
        output_names=output_names,  # 为输出节点设置别名
        # 如果不设置dynamic_axes,那么对于输入形状为[4, 3, 224, 224],在以后使用onnx进行推理时也必须输入[4, 3, 224, 224]
        # 下面设置了输入的第0维是动态的,以后推理时batch_size的大小可以是其他动态值
        dynamic_axes={
            # a dictionary to specify dynamic axes of input/output
            # each key must also be provided in input_names or output_names
            "input_data": {0: "batch_size"},
            "output_data": {0: "batch_size"}
        })
    return model_path


def onnx_check(model_path):
    """
    验证导出的模型格式时候正确
    :param model_path:
    :return:
    """
    onnx_model = onnx.load(model_path)
    onnx.checker.check_model(onnx_model)
    print(onnx.helper.printable_graph(onnx_model.graph))


def onnx_inference(model_path):
    """
    模型推理
    :param model_path: 
    :return: 
    """
    # 使用onnxruntime-gpu在GPU上进行推理
    session = onnxruntime.InferenceSession(model_path,
                                           providers=[
                                               ("CUDAExecutionProvider", {  # 使用GPU推理
                                                   "device_id": 0,
                                                   "arena_extend_strategy": "kNextPowerOfTwo",
                                                   "gpu_mem_limit": 4 * 1024 * 1024 * 1024,
                                                   "cudnn_conv_algo_search": "EXHAUSTIVE",
                                                   "do_copy_in_default_stream": True,
                                                   # "cudnn_conv_use_max_workspace": "1"    # 在初始化阶段需要占用好几G的显存
                                               }),
                                               "CPUExecutionProvider"       # 使用CPU推理
                                           ])
    # session = onnxruntime.InferenceSession(model_path)
    data = np.random.randn(2, 3, 224, 224).astype(np.float32)

    # 获取模型原始输入的字段名称
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    print("input name: {}".format(input_name))

    # 以字典方式将数据输入到模型中
    outputs = session.run([output_name], {input_name: data})
    print(outputs)


if __name__ == '__main__':
    model_path = pytorch_2_onnx()

    onnx_check(model_path)

    onnx_inference(model_path)

来源:胖胖大海

物联沃分享整理
物联沃-IOTWORD物联网 » pytorch.onnx.export方法参数详解,以及onnxruntime-gpu推理性能测试

发表评论