🤗 引言

它是什么。 又一个 Ultralytics 的 YOLOv5 实现。yolort 致力于使目标检测任务的训练与推理更加无缝地结合在一起。目前，yolort 采用了与官方 YOLOv5 相同的模型结构。显著的区别在于，我们引入了动态形状机制，在此框架下，可以将预处理（letterbox）和后处理（nms）嵌入到模型图中，从而简化部署策略。从这个意义上讲，yolort 使得在 LibTorch、ONNX Runtime、TVM、TensorRT 等平台上更轻松、更友好地部署目标检测成为可能。

关于代码。 遵循 detr 的设计原则：

目标检测不应比分类更困难，也不应需要复杂的库来进行训练和推理。

yolort 的实现和实验都非常简单。你喜欢 torchvision 中的 faster-rcnn、retinanet 或 detr 的实现吗？你是否喜欢 yolov5？那你一定会爱上 yolort！

🆕 最新内容

• 2021年12月27日：新增 TensorRT C++ 接口示例。感谢 Shiquan。
• 2021年12月25日：支持导出为 TensorRT 模型，并使用 TensorRT Python 接口进行推理。
• 2021年9月24日：新增 ONNX Runtime C++ 接口示例。感谢 Fidan。
• 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 安装

  从 PyPI 简单安装：

  pip install -U yolort
  

  或者从源码安装：

  # 克隆 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")
  # 对多张图片进行推理
  predictions = model.predict(["bus.jpg", "zidane.jpg"])
  

通过 torch.hub 加载

这些模型也可以通过 torch hub 获取。要加载带有预训练权重的 yolov5s，只需执行以下操作：

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

从官方 yolov5 加载检查点

以下是加载由 ultralytics/yolov5 训练的检查点权重的接口。有关我们如何与 yolov5 共享以及如何与其 不同 的更多信息，请参阅我们的文档。

from yolort.models import YOLOv5

# 从 https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt 下载检查点
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 部署 yolort 的流程。

from yolort.runtime import PredictorORT

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

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

请查看这篇教程，了解如何使用 yolort 的 ONNX 模型转换及 ONNX Runtime 推理。此外，您还可以使用示例来调用 ONNX Runtime 的 C++ 接口。

在 TensorRT 后端上进行推理

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 模型转换为 TensorRT 格式，并使用 Python 接口。如果您想使用 C++ 接口，请参阅此示例。

🎨 模型图可视化

现在，yolort 可以直接绘制模型图。请查看我们的教程，了解如何使用和可视化模型图。

👋 贡献

我们非常欢迎您的参与！请参阅我们的贡献指南，了解如何开始以及如何帮助我们。感谢所有贡献者！如果您喜欢这个项目，请考虑给本仓库加个 ⭐，这是对我们最简单的支持方式。

📖 引用 yolort

如果您在论文或其他出版物中使用了 yolort，请使用以下 BibTeX 条目进行引用。

@Misc{yolort2021,
  author =       {Zhiqiang Wang、Song Lin、Shiquan Yu、Wei Zeng、Fidan Kharrasov},
  title =        {YOLORT：用于专用加速器上的目标检测运行时栈},
  howpublished = {\url{https://github.com/zhiqwang/yolort}},
  year =         {2021}
}

🎓 致谢

• yolov5 的实现借鉴了 ultralytics 的代码。
• 本仓库的架构设计及部分代码参考了 torchvision。

YOLOv5 Runtime Stack (yolort) 快速上手指南

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 3.6 - 3.10
• 核心依赖:
  • PyTorch >= 1.8.0
  • torchvision >= 0.9.0
• 可选依赖 (如需在 COCO 数据集上进行评估): pycocotools

提示: 请优先参考 PyTorch 官方安装指南 安装适合您 CUDA 版本的 PyTorch。国内用户可使用清华源加速安装：

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

您可以选择通过 PyPI 直接安装，或从源码安装以获取最新功能。

方式一：通过 PyPI 安装（推荐）

pip install -U yolort

国内加速:

pip install -U yolort -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

# 克隆仓库
git clone https://github.com/zhiqwang/yolort.git
cd yolort
# 以可编辑模式安装
pip install -e .

安装额外依赖 (可选)

如果需要运行 COCO 评估脚本：

pip install -U 'git+https://github.com/ppwwyyxx/cocoapi.git#subdirectory=PythonAPI'

基本使用

yolort 提供了极简的 API，支持直接加载预训练模型并进行推理。

1. 快速推理示例

以下代码演示了如何加载 yolov5s 模型并对单张或多张图片进行目标检测：

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

2. 通过 Torch Hub 加载

无需克隆仓库，直接使用 torch.hub 加载模型：

import torch

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

3. 加载 Ultralytics 官方权重

yolort 兼容 Ultralytics 官方训练的 .pt 权重文件：

from yolort.models import YOLOv5

# 指定官方权重路径 (需先下载 yolov5s.pt)
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)

4. 部署后端示例

yolort 的核心优势在于轻松部署到不同后端。

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 推理:

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