Yolov5转RKNN模型前的准备-以RK1126为例

pt模型转onnx模型代码

转onnx模型的代码在yolov5工程下models文件夹下–export.py。

修改yolo.py文件

转onnx模型前要先修改models文件夹下的yolo.py，如下图所示，注释红色框内的代码（训练模型时要解开注释），用绿色框中的代码代替（训练时注释掉）。

z = []
for i in range(self.nl):
    x[i] = self.m[i](x[i])
return x

在终端激活torch环境（训练yolov5的环境）cd到yolov5工程目录下输入 python models/export.py --weights ./best.pt（best.pt是yolo训练完的模型）。出现ONNX export success, saved as ./best.onnx，就算是成功了。

部分模型结构如下图，需要查看onnx模型结构的小伙伴直接去https://netron.app就可以了。

未屏蔽第二张图红色框代码，转成onnx模型结构。

今天就记录到这里了，后面转rknn模型以及图像或视频的推理会有后续的。