深度学习—权重文件格式.pt转.tflite

内容概要

在智能车比赛中，我使用基于pytorch框架下yolov5的改进算法yolov5_mobileNet训练得到的.pt格式的权重文件无法搭载在比赛所使用平台上，需要将.pt格式文件转换成.tflite格式，中间遇到不少问题，下面从环境配置到代码实现做一个总结

整体架构流程

1.将.pt文件转成.tflite文件的步骤

大概了解到两种方法，第一种是直接使用yolov5源码中的export.py文件把.pt格式转换成.tflite格式，但我一直没运行成功，不清楚原因；第二种是将.pt文件转换成一种中间表示形式.onnx后再转换成.tflite格式，具体流程是：.pt ------ .onnx ------ .pb ------ .tflite

2.先对上面几种文件格式简单介绍

• .pt ：.pt是PyTorch模型文件的扩展名，使用.pt文件格式来保存、加载和部署训练好的神经网络模型。这些模型包含了网络的结构、权重以及其他相关的训练参数。用户可以使用PyTorch库来加载.pt文件，并使用这些模型进行推理、预测或微调
• .onnx ：ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式。旨在解决不同深度学习框架之间的模型互操作性问题
• .pb ：.pb文件是用于存储和序列化TensorFlow模型的二进制文件格式，其中包含了模型的结构、权重参数和计算图等信息。是一种用于跨平台和跨语言的模型表示方式，适用于在各种环境和设备上加载和运行TensorFlow模型
• .tflite ：.tflite文件是TensorFlow Lite模型的文件格式，它是一种轻量级的模型表示方式。TensorFlow Lite是TensorFlow的衍生版本，为移动设备和嵌入式设备提供了一种高效运行机器学习模型的解决方案。.tflite文件经过模型量化（Quantization）和优化，可以在资源受限的设备上高效地运行。该文件格式使用FlatBuffers库进行编码，可以在移动设备上快速加载并执行模型。通过将模型转换为.tflite格式，可以在移动应用或边缘设备上进行实时推理，而无需依赖较大而复杂的TensorFlow模型

3.环境配置

1.安装tensorflow（我使用的是tensorflow2.8.0版本）

pip install tensorflow==2.8.0

2.安装tensorflow-addons对应版本

其他版本可以去这里找: https://gitcode.com/tensorflow/addons/overview

3.根据tensorflow版本安装ONNX和ONNX-TensorFlow

其他版本可以去这里找: https://github.com/onnx/onnx-tensorflow/tree/main

4.根据tensorflow版本安装tensorflow_probability和JAX

这一