特别说明:参考官方开源的yolov8代码、瑞芯微官方文档、地平线的官方文档,如有侵权告知删,谢谢。
模型和完整仿真测试代码,放在github上参考链接 模型和代码。
之前写了yolov8、yolov8seg、yolov8obb 的 DFL 放在模型中和放在后处理中的两种不同部署方法(推荐放在后处理中);而yolov8pose只写了一篇放在模型中,有网友希望写一篇放在后处理中的,yolov8pose的DFL放在后处理中的博客来了。
1 模型和训练
老规矩,训练不涉及,训练代码参考官方开源的yolov8训练代码。
2 导出 yolov8 onnx
修改两处、增加保存onnx代码,一共修改三个地方。
修改第一处:增加以下几行代码
python
# 导出 onnx 增加
y = []
for i in range(self.nl):
t1 = self.cv2[i](x[i])
t2 = self.cv3[i](x[i])
y.append(t1)
y.append(t2)
return y
第二处:增加以下几行代码
python
ps = []
for i in range(self.nl):
ps.append(self.cv4[i](x[i]))
x = self.detect(self, x)
return x, ps
增加保存onnx模型代码
python
print("=========== onnx =========== ")
dummy_input = torch.randn(1, 3, 640, 640)
input_names = ["data"]
output_names = ["reg1", "cls1", "reg2", "cls2", "reg3", "cls3", "ps1", "ps2", "ps3"]
torch.onnx.export(self.model, dummy_input, "./yolov8n_pose.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=11)
print("======================== convert onnx Finished! .... ")
修改完以上几个地方,运行以下两行。
python
from ultralytics import YOLO
model = YOLO('./weights/yolov8n-pose.pt')
3 yolov8n-pose onnx 测试效果
4 时耗
将DFL放后处理中,模型和后处理时耗,使用芯片rk3588,模型输入分辨率640x640。部署rk3588板端代码
以下截图是将DFL放在模型中的时耗(模型推理时会长一些,但后处理会稍微快一点),将DFL放在模型中部署【参考链接】