如何将原始yolo权重转换为TensorRT模型

Question

我开发了yolov4微小模型的改进版本。我想将这个开发的模型转换为TensorRT模型，但在参考了附加的URL后，我发现我只能转换原始的v4微小模型。

我的问题是，其他人是如何将他们的原始模型转换为TensorRT的？

提前谢谢。URL

Answer 1

我知道您有一个自己训练过的自定义模型，并且希望将其转换为TensorRT。

有很多方法可以将模型转换为TensorRT。这个过程取决于你的模型是哪种格式，但这里有一种适用于所有格式：

1. 将模型转换为ONNX格式
2. 使用trtexec将模型从ONNX转换为TensorRT

详细步骤

我想你的模型是Pytorch格式的。至少您链接的存储库中的train.py会将模型保存为该格式。您可以使用tf2onnx将其转换为ONNX。

请注意，tf2onnx建议使用Python 3.7。如果您使用的是另一个版本的Python，您可以在这里安装它，并使用conda或venv创建一个虚拟环境。

然后，安装tf2onnx:

pip install git+https://github.com/onnx/tensorflow-onnx

将您的模型从保存的模型转换为ONNX

python3 -m tf2onnx.convert --saved-model ./model --output model.onnx 

如果您的模型使用其他tf格式，请参阅tf2onnx的自述文件以获得帮助。

如果您还没有安装TensorRT及其依赖项，请使用本指南进行安装。或者，您也可以使用Nvidia Containers(NGC(。

安装TensorRT后，您可以执行此命令来使用fp16精度转换模型。

/usr/src/tensorrt/bin/trtexec --onnx=model.onnx --saveEngine=model.engine --fp16 --workspace=3000 --buildOnly

您可以通过运行来检查所有CLI参数

/usr/src/tensorrt/bin/trtexec --help

Answer 2

对于YOLO v3-v5，您可以使用手动解析cfg和weight文件的项目，手动构建和保存TensorRT的引擎文件。例如enazoe/yolo tensorrt。我在Windows/Linux x86/Nvidia Jetson上使用这个内置多目标跟踪器作为快速对象检测器。

在这种情况下，您不需要安装trtexec和NVidia的其他软件。