yolort

🤗 Introduction

What it is. The YOLOv5 Runtime Stack, or yolort, is an implementation of the well-known YOLOv5 model created by Ultralytics. This project aims to make object detection tasks easier by integrating training and inference into a seamless experience. Unlike some other implementations, yolort uses a dynamic shape mechanism, allowing both pre-processing (letterbox) and post-processing (NMS - Non-Maximum Suppression) to be embedded directly into the model's graph. This simplification makes deploying object detection models across platforms like LibTorch, ONNX Runtime, TVM, and TensorRT much more accessible.

About the code. Following the design principle of DETR, yolort underscores the premise that object detection should not be more complicated than classification and should not necessitate complex libraries for training and inference. If you enjoy working with the simpler implementations of models like faster-rcnn, retinanet, or detr in PyTorch's torchvision, you'll likely appreciate yolort.

🆕 What's New

The yolort project continuously evolves with regular updates:

• Dec. 27, 2021: Integration of a TensorRT C++ interface example.
• Dec. 25, 2021: Support for exporting to TensorRT and associated Python interface for inference.
• Sep. 24, 2021: Addition of an ONNX Runtime C++ interface example.
• Feb. 5, 2021: Incorporation of TVM compile and inference notebooks.
• Nov. 21, 2020: Introduction of graph visualization tools.
• Nov. 17, 2020: Support for exporting to ONNX and inference using the ONNX Runtime Python interface.
• Nov. 16, 2020: Refactoring of YOLO modules with support for dynamic shape/batch inference.
• Nov. 4, 2020: Introduction of a LibTorch C++ inference example.
• Oct. 8, 2020: Support for exporting to the TorchScript model.

🛠️ Usage

The simplicity and minimal package dependencies of yolort make it easy to use without any extra compiled components.

Installation and Inference Examples

Begin by installing PyTorch version 1.8.0 or higher and torchvision version 0.9.0 or higher. Then, you can install yolort via pip from PyPI or from the source:

pip install -U yolort
# OR
git clone https://github.com/zhiqwang/yolort.git
cd yolort
pip install -e .

For evaluation on COCO, install pycocotools:

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

To perform inference:

from yolort.models import yolov5s

# Load model
model = yolov5s(pretrained=True, score_thresh=0.45)
model.eval()

# Perform inference on an image file
predictions = model.predict("bus.jpg")
# or on a list of image files
predictions = model.predict(["bus.jpg", "zidane.jpg"])

Loading via torch.hub

Models are available through torch hub. Load the yolov5s with pretrained weights as follows:

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

Loading checkpoint from official yolov5

To load checkpoint weights trained with ultralytics/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)

🚀 Deployment

yolort supports deployment on various backends:

Inference on LibTorch backend

The tutorial shows how to convert models into torchscript and perform inference with a serialized torchscript model in C++.

Inference on ONNX Runtime backend

Deploying with ONNX Runtime is made straightforward with yolort:

from yolort.runtime import PredictorORT

engine_path = "yolov5n6.onnx"
y_runtime = PredictorORT(engine_path, device="cpu")

predictions = y_runtime.predict("bus.jpg")

Follow the tutorial for guidance on using yolort for ONNX model conversion and inference. Additionally, a C++ example is available for better understanding.

Inference on TensorRT backend

Deploying with TensorRT is equally intuitive:

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

The tutorial provides further details on converting yolort models to TensorRT and using them with Python. A relevant C++ example is also available.

🎨 Model Graph Visualization

yolort provides tools for visualizing the model graph. The tutorial explains how to use and create visualizations.

👋 Contributing

Community contributions are encouraged! Please refer to the Contributing Guide for instructions on getting involved. A star on the repository is the simplest way to acknowledge and support the project.

📖 Citing yolort

If yolort contributes to your research or publications, consider citing it with the following BibTeX entry:

@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}
}

🎓 Acknowledgement

The yolort project took inspiration from Ultralytics's yolov5 and borrowed architectural design elements from PyTorch's torchvision.