date
Aug 7, 2025
slug
2025-08-07-the-onnx-export-and-deployment-of-pytorch-model
status
Published
tags
CNN
ONNX
summary
本文详细总结了基于Pytorch框架开发和训练好的模型导出为通用的onnx文件，并且通过onnxruntime进行部署和Python/C++访问的流程和步骤。
type
Post
category
AI
AI summary
本文详细总结了基于Pytorch框架开发和训练好的模型导出为通用的onnx文件，并且通过onnxruntime进行部署和Python/C++访问的流程和步骤。
本文要解决的问题是：在Pytorch中开发和训练好的网络框架及其参数，如何导出为可以部署的onnx文件？如何在其他的环境中部署这个onnx文件？ 本文以[[基于Pytorch实现手写数字识别的卷积神经网络]]一文中所训练的卷积神经网络为例，来说明这个完整的过程。
Pytorch模型导出为pt/pth文件
在Pytorch框架中，已经训练完成或者训练中的模型及其参数可以保存pt或者pth文件，这样后续就可以直接通过加载这些文件快速恢复出来模型及其训练后的参数状态，继续进行训练或者直接用于新数据的推理。
一般而言，使用pt后缀名的文件用于仅保存模型的参数，而pth后缀名文件则同时保存了模型的结构和参数。因此前者保存的文件更小，但在下次加载前需要提前构建模型的框架结构，因为pt文件里面并没有包含模型结构；而后者在使用上更加方便，内部同时包含了模型结构和参数，使用时直接加载即可，但是文件体积更大。
当然，在 PyTorch 中， pt和pth这两种文件后缀没有官方严格的区分和定义，它们更多是开发者社区约定俗成的用法。也就是说，你完全可以用 .pt 保存完整模型，或用 .pth 仅保存参数。
以下是基于Pytorch保存训练模型及其参数，并且重新加载模型参数到新的模型对象的代码示例：
模型的ONNX导出
相比于Pytorch的模型参数导出的pt/pth文件格式，ONNX文件格式是模型部署以及在不同模型开发框架之间共享和转换支持上更为普遍的文件格式。因此也可以把训练的模型及其参数导出为ONNX文件，用于进行模型部署以及转换到其他的框架下。
Pytorch模型导出onnx文件并通过ONNXRuntime类似进行新样本推理的前提条件，是需要安装几个依赖库：
把已经训练好的模型及其参数导出为Onnx文件的代码流程如下：
如以上代码中input_sample变量所表明的，在进行ONNX文件导出时需要明确的指定后续进行推理的数据输入格式。以上代码执行后，在当前目录下就会生成一个model.onnx文件，后续就可以基于这个文件进行Onnx部署以及在不同框架下转换和使用。
ONNX文件检查
以上所导出的ONNX文件，可以使用以下代码检测该文件的合法性：
onnx.checker.check_model接口用于验证ONNX模型文件的有效性，确保它符合ONNX规范，可以被ONNX兼容的推理引擎正确加载和执行。这是在使用ONNX模型进行推理前的一个常见检查步骤，有助于提高代码的健壮性。
此外，也可以直接使用netron.app工具来对以上onnx文件进行检查，该工具可以以图形的方式输出非常详细的模型及其参数信息，可以非常直观的看到该文件中所描述的模型结构及其当前各层的训练参数信息：
通过ONNXRuntime部署模型并测试
Onnx文件中所保存的模型结构和参数是模型的静态表示，不能够直接调用PT格式导致后的eval()和predict()方法。当基于Onnx文件加载模型的结构和参数以后，需要通过ONNXRuntime引擎来执行新样本数据的推理。因此要首先安装ONNXRuntime包才能支持对ONNX文件的部署和推理。
如前所述，在Pytorch框架下使用ONNXRuntime导入Onnx文件并进行新样本的推理，需要通过pip安装onnxruntime和protobuf两个包。
以前文所导出的模型Onnx文件为例，从MNIST数据集中的某个样本进行OnnxRuntime的推理测试代码如下：
可以看到，除了使用torchvision.datasets和torchvision.transforms来实现对数据集的访问以外（其实完全可以采用其他方式来准备输入该模型进行推理的数据样本，只要数据格式符合ONNX input sample的要求即可），以上代码对于新样本数据的推理仅仅依赖于numpy和onnxruntime，完全不依赖于开发该模型的pytorch框架。