访问官网
Github
Huggingface
文档
论文YOLOv5 Runtime Stack 项目介绍
🤗 项目简介
YOLOv5 Runtime Stack（简称 yolort）是Ultralytics公司YOLOv5的又一个实现版本。yolort旨在让目标检测任务的训练和推理环节更加无缝衔接。目前，yolort采用与官方YOLOv5相同的模型结构，但不同之处在于它采用了动态形状机制，将预处理（如letterbox）和后处理（如nms）嵌入到模型图中，这样可以简化部署策略。因此，yolort可以更轻松地在LibTorch、ONNX Runtime、TVM、TensorRT等平台上部署目标检测任务。
yolort的代码基于detr的设计原则，简单易用，训练和推理不再需要复杂的库。如果你喜欢torchvision的faster-rcnn、retinanet或detr，或者是YOLOv5，那么你一定会喜欢yolort。
🆕 更新动态
- 2021年12月27日：新增TensorRT C++接口示例。
- 2021年12月25日：支持导出到TensorRT，并新增TensorRT Python接口推理。
- 2021年9月24日：新增ONNX Runtime C++接口示例。
- 2021年2月5日：新增TVM编译与推理笔记。
- 2020年11月21日：新增图形可视化工具。
- 2020年11月17日：支持导出到ONNX，并新增ONNX Runtime Python接口推理。
- 2020年11月16日：改进YOLO模块，支持动态形状/批量推理。
- 2020年11月4日：新增LibTorch C++推理示例。
- 2020年10月8日：支持导出到TorchScript模型。
🛠️ 使用方法
yolort没有额外的编译组件，并且软件包依赖关系很少，因此代码非常易于使用。
安装和推理示例
-
首先，根据官方说明安装PyTorch 1.8.0+和torchvision 0.9.0+。
-
通过pip安装：
pip install -U yolort -
或者从源代码安装：
git clone https://github.com/zhiqwang/yolort.git cd yolort pip install -e . -
安装pycocotools（用于COCO评估）：
pip install -U 'git+https://github.com/ppwwyyxx/cocoapi.git#subdirectory=PythonAPI' -
图像文件推理示例：
from yolort.models import yolov5s model = yolov5s(pretrained=True, score_thresh=0.45) model.eval() predictions = model.predict("bus.jpg")
使用torch.hub加载
通过torch hub加载预训练的yolov5s模型：
model = torch.hub.load("zhiqwang/yolort:main", "yolov5s", pretrained=True)
从YOLOv5官方加载检查点
接口如下：
from yolort.models import YOLOv5
ckpt_path_from_ultralytics = "yolov5s.pt"
model = YOLOv5.load_from_yolov5(ckpt_path_from_ultralytics, score_thresh=0.25)
model.eval()
img_path = "test/assets/bus.jpg"
predictions = model.predict(img_path)
🚀 部署
LibTorch后端推理
有一个教程演示如何将模型转换为torchscript，并提供一个C++示例以实现使用序列化的torchscript模型进行推理。
ONNX Runtime后端推理
以下是使用ONNX Runtime进行部署的管道：
from yolort.runtime import PredictorORT
engine_path = "yolov5n6.onnx"
y_runtime = PredictorORT(engine_path, device="cpu")
predictions = y_runtime.predict("bus.jpg")
TensorRT后端推理
以下是TensorRT部署的简单示例：
import torch
from yolort.runtime import PredictorTRT
engine_path = "yolov5n6.engine"
device = torch.device("cuda")
y_runtime = PredictorTRT(engine_path, device=device)
predictions = y_runtime.predict("bus.jpg")
🎨 模型图形可视化
yolort可以直接绘制模型图形，用户可以参考教程了解如何使用及可视化模型图。
👋 撰稿与社区贡献
欢迎社区贡献者提供帮助。用户可以参考我们的贡献指南开始合作。如果喜欢该项目，请考虑为该仓库加星，以支持我们。
📖 引用yolort
如果您在您的出版物中使用yolort，请使用以下BibTeX条目进行引用。
🎓 鸣谢
yolov5的实现借用了ultralytics的代码。- 本仓库的架构设计及部分代码来自于torchvision。
