YOLORT: 适用于多种推理加速器的YOLOv5运行时栈

yolort

YOLORT:面向多种推理加速器的YOLOv5运行时栈

YOLORT是Ultralytics公司YOLOv5的另一种实现,旨在让目标检测的训练和推理更加无缝集成。作为一个强大的目标检测框架,YOLORT在保持YOLOv5高性能的同时,通过创新的设计使得模型部署变得更加简单和灵活。

🚀 YOLORT的创新之处

YOLORT采用了与官方YOLOv5相同的模型结构,但有几个显著的创新点:

动态形状机制: YOLORT引入了动态形状机制,这使得模型能够更好地适应不同尺寸的输入图像。
预处理和后处理集成: 将letterbox预处理和非极大值抑制(NMS)后处理嵌入到模型图中,简化了部署流程。
多加速器支持: 可以轻松部署到LibTorch、ONNX Runtime、TVM、TensorRT等多种推理加速器上。

这些创新大大简化了目标检测模型的部署策略,使得YOLORT在各种硬件平台上的应用变得更加容易和友好。

💡 设计理念

YOLORT的设计遵循了Facebook Research的DETR项目的设计原则:

目标检测的难度不应高于分类,也不应需要复杂的库来进行训练和推理。

基于这一理念,YOLORT的实现非常简洁,易于使用和实验。如果你喜欢torchvision中faster-rcnn、retinanet或detr的实现,或者你喜欢YOLOv5,那么你一定会喜欢上YOLORT!

🛠️ 使用方法

YOLORT的使用非常简单,没有额外的编译组件,依赖项也很少。以下是一些基本用法:

安装:
```
pip install yolort
```

加载预训练模型并进行推理:

from yolort.models import yolov5s

# 加载模型
model = yolov5s(pretrained=True, score_thresh=0.45)
model.eval()

# 对单张图片进行推理
predictions = model.predict("bus.jpg")

# 对多张图片进行推理
predictions = model.predict(["bus.jpg", "zidane.jpg"])

通过torch.hub加载:

model = torch.hub.load("zhiqwang/yolort:main", "yolov5s", pretrained=True)

🚀 部署

YOLORT提供了多种部署选项,使其能够在不同的推理后端上运行:

LibTorch后端: YOLORT提供了将模型转换为TorchScript的教程,以及使用序列化TorchScript模型进行C++推理的示例。

ONNX Runtime后端:

from yolort.runtime import PredictorORT

# 加载序列化的ONNX模型
engine_path = "yolov5n6.onnx"
y_runtime = PredictorORT(engine_path, device="cpu")

# 对图片进行推理
predictions = y_runtime.predict("bus.jpg")

TensorRT后端:

import torch
from yolort.runtime import PredictorTRT

# 加载序列化的TensorRT引擎
engine_path = "yolov5n6.engine"
device = torch.device("cuda")
y_runtime = PredictorTRT(engine_path, device=device)

# 对图片进行推理
predictions = y_runtime.predict("bus.jpg")

这些部署选项使YOLORT能够在各种硬件平台上高效运行,满足不同场景的需求。

🎨 模型图可视化

YOLORT还提供了模型图的可视化功能,让用户能够直观地了解模型结构。通过简单的API调用,就可以生成清晰的模型图示:

YOLO model visualize

这一功能对于理解模型结构、调试和优化都非常有帮助。

👥 社区贡献

YOLORT是一个开源项目,欢迎社区成员的贡献。如果你对这个项目感兴趣,可以通过以下方式参与:

提交Issue报告bug或提出新功能建议
提交Pull Request贡献代码
改进文档
分享使用经验

社区的支持和贡献是项目持续发展的动力。如果你觉得YOLORT对你有帮助,别忘了给项目的GitHub仓库点个星⭐️!

📚 引用

如果你在研究中使用了YOLORT,请使用以下BibTeX条目进行引用:

@Misc{yolort2021,
  author =       {Zhiqiang Wang and Song Lin and Shiquan Yu and Wei Zeng and Fidan Kharrasov},
 title =        {YOLORT: A runtime stack for object detection on specialized accelerators},
  howpublished = {\url{https://github.com/zhiqwang/yolort}},
  year =         {2021}
}