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)