yolort

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")