PyTorch 训练之后的网络模型.pth转.onnx文件并对图像进行预测

目录

模型文件格式pth和onnx

pth文件转ONNX格式的优势

---

B站视频讲解https://www.bilibili.com/video/BV1E2vMeTEkr?p=8&vd_source=b2eaaddb2c69bf42517a2553af8444ab

GitHub源码https://github.com/KeepTryingTo/DeepLearning/tree/main/Classification/pth2onnx

本文主要是基于图像分类 算法的实现，将训练得到的.pth权重文件转换为ONNX文件，关于ONNX模型文件介绍如下。

图像分类

PyTorch实现一个简单的图像分类（代码详细）

模型文件格式pth和onnx

.pth文件	.onnx文件
使用PyTorch深度学习框架训练模型之后一般保存的格式为.pth文件（权重文件）或者以字典的形式保存模型文件（一般包含模型权重，优化器以及当前训练迭代步数epoch等信息），不管怎么样，PyTorch框架训练之后保存的模型文件主要包含网络结构，层参数，一般用于后期的微调或者推理。	PyTorch框架训练之后得到的模型如果要想跨平台或者部署落地的话，需要转换为ONNX模型格式，可以用于C++，Java，JavaScript等语言环境中，并且从pth转换为ONNX格式之后，可以兼容不同的框架，可以用于嵌入式设备，移动设备等计算资源有限的硬件中。

注：大家目前只需要知道pth文件转换为ONNX格式有以下好处即可，关于ONNX的其他使用后期慢慢学习中体会。

pth文件转ONNX格式的优势

• 跨平台兼容性
  • 通用格式：ONNX可以作为一种通用的中间表示，可以再各种硬件和平台上进行加速；
  • 支持多种设备：ONNX模型文件可以在CPU,GPU,FPGA等多种设备上运行。
• 性能优化
  • 高级优化技术：ONNX支持一系列的优化技术，比如精度剪枝（从F32->F16），通道剪枝（减少那些贡献少的通道），模型量化（模型使用更少bit表示）和压缩等，主要是让模型更加的轻量化。
  • 推理加速：既然对其模型进行了优化，自然也加速了模型的推理速度。并且使用ONNX Runtime推理引擎进一步优化ONNX模型的推理性能。
• 模型部署的灵活性
  • 易于迁移：转换为ONNX可以很容易将其迁移到支持ONNX的各种深度学习框架中，比如Tensorflow,MxNet等。
  • 多平台支持：和上面讲的差不多，可以在多种硬件上运行和各种操作系统环境中部署，提高了模型部署的灵活性和便捷性。
• 模型的可靠性和稳定性都提升
  • 模型验证：ONNX对模型进行一系列的检查和验证，减少了在验证过程中的错误和故障。
  • 标准化流程：ONNX提供了标准化的模型转换和部署流程，可以降低由不同框架和平台之间的差异而导致的风险。
• 易于集成和扩展
  • 丰富的生态系统：由于很多工具和库都支持ONNX，那么自然而然的ONNX的应用也越来越广泛，为模型的集成和扩展提供了丰富的选择；
  • 社区支持：拥有庞大的社区，对于ONNX的支持和推广具有很大的影响。

思路：

• 将,pth文件转换为ONNX模型文件；

注：关于opset_version更多算子，请看https://github.com/onnx/onnx/blob/main/docs/Operators.md 

• 加载ONNX模型文件和使用ONNXRuntime推理引擎加速模型推理；

• 输出预测结果以及注意地方。

图像分类模型.pth转ONNX模型文件