YOLOrt 项目使用教程

YOLOrt 项目使用教程

yolort yolort is a runtime stack for yolov5 on specialized accelerators such as tensorrt, libtorch, onnxruntime, tvm and ncnn. 项目地址: https://gitcode.com/gh_mirrors/yo/yolort

1. 项目介绍

YOLOrt 是一个为 YOLOv5 设计的运行时堆栈，旨在使对象检测任务的训练和推理更加无缝集成。YOLOrt 采用了与官方 YOLOv5 相同的模型结构，但引入了动态形状机制，并将预处理（如 letterbox）和后处理（如 NMS）嵌入到模型图中，从而简化了部署策略。YOLOrt 支持在多种专用加速器上运行，如 TensorRT、LibTorch、ONNX Runtime、TVM 和 NCNN。

2. 项目快速启动

安装

首先，确保你已经安装了 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 .

推理示例

加载模型并进行推理

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 加载模型

import torch

# 通过 torch.hub 加载预训练模型
model = torch.hub.load("zhiqwang/yolort:main", "yolov5s", pretrained=True)

3. 应用案例和最佳实践

案例1：在 LibTorch 后端进行推理

YOLOrt 提供了一个教程，展示了如何将模型转换为 TorchScript，并提供了一个 C++ 示例，展示了如何使用序列化的 TorchScript 模型进行推理。

案例2：在 ONNX Runtime 后端进行推理

YOLOrt 提供了一个部署管道，用于在 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")

案例3：在 TensorRT 后端进行推理

YOLOrt 还提供了一个用于 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")

4. 典型生态项目

TensorRT

TensorRT 是一个高性能的深度学习推理库，YOLOrt 提供了对 TensorRT 的支持，使得模型可以在 NVIDIA GPU 上高效运行。

ONNX Runtime

ONNX Runtime 是一个跨平台的推理引擎，YOLOrt 支持将模型导出为 ONNX 格式，并在 ONNX Runtime 上进行推理。

LibTorch

LibTorch 是 PyTorch 的 C++ 前端，YOLOrt 提供了将模型转换为 TorchScript 的教程，并提供了 C++ 推理示例。

TVM

TVM 是一个开源的深度学习编译器，YOLOrt 支持将模型编译为 TVM 格式，并在 TVM 上进行推理。

NCNN

NCNN 是一个为移动平台设计的深度学习推理框架，YOLOrt 支持在 NCNN 上进行推理。

yolort yolort is a runtime stack for yolov5 on specialized accelerators such as tensorrt, libtorch, onnxruntime, tvm and ncnn. 项目地址: https://gitcode.com/gh_mirrors/yo/yolort