ONNXRuntime学习笔记(三)

接上一篇完成的pytorch模型训练结果，模型结构为ResNet18+fc，参数量约为11M，最终测试集Acc达到94.83%。接下来有分两个部分：导出onnx和使用onnxruntime推理。

一、pytorch导出onnx

直接放函数吧，这部分我是直接放在test.py里面的，直接从dataloader中拿到一个batch的数据走一遍推理即可。

def export_onnx(net, testloader, output_file):
    net.eval()
    with torch.no_grad():
        for data in testloader: 
            images, labels = data

            torch.onnx.export(net, 
                            (images), 
                            output_file,
                            training=False,
                            do_constant_folding=True,
                            input_names=["img"], 
                            output_names=["output"],
                  dynamic_axes={"img": {0: "b"},"output": {0: "b"}}
                  )
            print("onnx export done!")
            break

上面函数中几个比较重要的参数：do_constant_folding是常量折叠，建议打开；输入张量通过一个tuple传入，并且最好指定每个输入和输出的名称，此外，为保证使用onnxruntime推理的时候batchsize可变，dynamic_axes的第一维需要像上述一样设置为动态的。如果是全卷积做分割的网络，类似的输入h和w也应该是动态的。

单独运行test.py计算测试集效果和平均相应时间,为方便比较，这里batch_size设置为1,结果为：

Test Acc is: 94.84%
Average response time cost: 8.703978610038757 ms

二、使用onnxruntime推理

这里我们使用gpu版本的onnxruntime库进行推理，其python包可直接pip install onnxruntime-gpu安装。onnxruntime推理代码和测试集推理代码很类似,如下：

import numpy as np
import onnxruntime as ort
import argparse, os
from lib import CIFARDataset

def onnxruntime_test(session, testloader):
    print("Start Testing!")
    input_name = session.get_inputs()[0].name
    correct = 0
    total = 0   # 计数归零（初始化）
    for data in testloader:
        images, labels = data
        images, labels = images.numpy(), labels.numpy()
        outputs = session.run(None, {input_name:images})
        predicted = np.argmax(outputs[0], axis=1)  # 取得分最高的那个类
        total += labels.shape[0]                        # 累加样本总数
        correct += (predicted == labels).sum()        # 累加预测正确的样本个数
    acc = correct / total
    print('ONNXRuntime Test Acc is: %.2f%%' % (100*acc))
            
if __name__ == '__main__':
    # 命令行参数解析
    parser = argparse.ArgumentParser("CNN backbone on cifar10")
    parser.add_argument('--onnx', default='./output/test_resnet18_10_autoaug/densenet_best.onnx')
    args = parser.parse_args()

    NUM_CLASS =10
    BATCH_SIZE = 1 # 批处理尺寸(batch_size)

    # 数据集迭代器
    data_path="./data"
    dataset = CIFARDataset(dataset_path=data_path, batchsize=BATCH_SIZE)
    _, testloader = dataset.get_cifar10_dataloader()

    # 构建session
    sess = ort.InferenceSession(args.onnx, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

    #onnxruntime推理
    import time
    start = time.time()
    onnxruntime_test(sess, testloader)
    end = time.time()
    print(f"Average response time cost: {1000*(end-start)/len(testloader.dataset)} ms")

使用onnxruntime加载导出的onnx模型，计算测试集效果和平均响应时间，结果为：

ONNXRuntime Test Acc is: 94.83%
Average response time cost: 3.1050602436065673 ms

三、小结

分析上面的pytorch和onnxruntime的测试结果可知，最终测试集效果是一致的，Acc分别为94.84%和94.83%，相当于10000个样本里面只有1个的预测结果不一致，这是可以接受范围内。但onnxruntime的效率更高，平均耗时只有3.1ms，比pytorch的8.7ms快了将近3倍。这在实际部署中的优势是非常明显的。目前Python端的结论比最初目标设定的50ms高很多，如果说需要进一步优化，两个方向：模型量化或并行化推理（拼batch或多线程）。下一篇再分析。