轻松学pytorch之使用onnx runtime实现pytorch模型部署

ONNX（Open Neural Network Exchange）是一种标准与开放的网络模型交换格式，直白点说就是tensorflow/pytorch/paddle等深度学习框架训练的模型都可以转换为ONNX格式，然后ONNX格式模型可以通过ONNX runtime组件实现模型的推理预测并加速，从而实现不基于原来框架的模型部署。ONNX runtime的github地址如下：

https://github.com/Microsoft/onnxruntime

支持操作系统与硬件平台如下：

Python开发环境下安装onnx runtime只需要一条命令行：

pip install onnxruntime

然后测试一下

这样就表示安装成功啦！

默认情况下，上述安装的onnxruntime只支持CPU推理，也就是说模型是运行的CPU版本，支持的数据类型为Numpy的Map或者数组或者List类型，模型默认在CPU上推理执行。测试代码如下：

def cpu_ort_demo():
    device_name = onnxruntime.get_device()
    print(device_name)
    session = onnxruntime.InferenceSession("D:/python/pytorch_tutorial/defect_detection/surface_defect_resnet18.onnx")
    image = cv.imread("D:/pytorch/enu_surface_defect/test/Pa_5.bmp")
    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
    x_input = img_transform(image).view(1, 3, 200, 200)
    ort_inputs = {session.get_inputs()[0].name: x_input.numpy()}
    ort_outs = session.run(None, ort_inputs)
    out_prob = ort_outs[0]
    label_id = np.argmax(out_prob)

    defect_txt = defect_labels[label_id]
    cv.putText(image, defect_txt, (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2, 8)
    cv.imshow("defect_detection", image)
    cv.waitKey(0)
    cv.destroyAllWindows()

运行结果如下（自定义缺陷分类）：

默认情况下，上述安装的onnxruntime只支持CPU推理，也就是说模型是运行的CPU版本，想要完成CUDA版本的推理，需要安装onnxruntime-gpu版本，安装的命令行如下：

pip install onnxruntime-gpu

使用GPU推理支持需要VC++与CUDA版本匹配支持，这个坑比较多，而且onnxruntime版本不同支持的CUDA版本也不一样。上面的代码输入改为CUDA支持版本如下：

def gpu_ort_demo():
    device_name = onnxruntime.get_device()
    print(device_name)
    session = onnxruntime.InferenceSession("D:/python/pytorch_tutorial/defect_detection/surface_defect_resnet18.onnx")
    image = cv.imread("D:/pytorch/enu_surface_defect/test/Pa_5.bmp")
    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
    x_input = img_transform(image).view(1, 3, 200, 200).numpy()
    print(x_input.shape)

    ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(x_input, 'cuda', 0)
    ortvalue.device_name()  # 'cuda'
    ortvalue.shape()  # shape of the numpy array X
    ortvalue.data_type()  # 'tensor(float)'
    ortvalue.is_tensor()  # 'True'
    np.array_equal(ortvalue.numpy(), x_input)  # 'True'

    ort_inputs = {session.get_inputs()[0].name: ortvalue}
    ort_outs = session.run(None, ort_inputs)
    out_prob = ort_outs[0]
    label_id = np.argmax(out_prob)

    defect_txt = defect_labels[label_id]
    cv.putText(image, defect_txt, (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2, 8)
    cv.imshow("defect_detection", image)
    cv.waitKey(0)
    cv.destroyAllWindows()

其中OrtValue会自动实现模型输入数据的CPU的内存数据到CUDA数据搬运，此外还可以通过会话的io_binding()方法实现对输入与输出数据的CUDA绑定支持。

OpenVINO2021.3开始支持onnxruntime方式加速推理，前提是需要自己编译，关于编译支持操作可以查看下面的链接：

https://docs.openvinotoolkit.org/2021.3/omz_tools_accuracy_checker_onnx_runtime_launcher.html