Skip to main content

Yet Another YOLOv5 and its Additional Runtime Stack

Project description

:hugs: Introduction

What it is. Yet another implementation of Ultralytics's YOLOv5. yolort aims to make the training and inference of the object detection integrate more seamlessly together. yolort now adopts the same model structure as the official YOLOv5. The significant difference is that we adopt the dynamic shape mechanism, and within this, we can embed both pre-processing (letterbox) and post-processing (nms) into the model graph, which simplifies the deployment strategy. In this sense, yolort makes it possible to be deployed more friendly on LibTorch, ONNXRuntime, TVM and so on.

About the code. Follow the design principle of detr:

object detection should not be more difficult than classification, and should not require complex libraries for training and inference.

yolort is very simple to implement and experiment with. You like the implementation of torchvision's faster-rcnn, retinanet or detr? You like yolov5? You love yolort!

YOLO inference demo

:new: What's New

  • Sep. 24, 2021. Add ONNXRuntime C++ interface example. Thanks to itsnine.
  • Feb. 5, 2021. Add TVM compile and inference notebooks.
  • Nov. 21, 2020. Add graph visualization tools.
  • Nov. 17, 2020. Support exporting to ONNX, and inferencing with ONNXRuntime Python interface.
  • Nov. 16, 2020. Refactor YOLO modules and support dynamic shape/batch inference.
  • Nov. 4, 2020. Add LibTorch C++ inference example.
  • Oct. 8, 2020. Support exporting to TorchScript model.

:hammer_and_wrench: Usage

There are no extra compiled components in yolort and package dependencies are minimal, so the code is very simple to use.

Installation and Inference Examples

  • Above all, follow the official instructions to install PyTorch 1.7.0+ and torchvision 0.8.1+

  • Installation via Pip

    Simple installation from PyPI

    pip install -U yolort
    

    Or from Source

    # clone yolort repository locally
    git clone https://github.com/zhiqwang/yolov5-rt-stack.git
    cd yolov5-rt-stack
    # install in editable mode
    pip install -e .
    
  • Install pycocotools (for evaluation on COCO):

    pip install -U 'git+https://github.com/ppwwyyxx/cocoapi.git#subdirectory=PythonAPI'
    
  • To read a source of image(s) and detect its objects :fire:

    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")
    # Perform inference on a list of image files
    predictions = model.predict(["bus.jpg", "zidane.jpg"])
    

Loading via torch.hub

The models are also available via torch hub, to load yolov5s with pretrained weights simply do:

model = torch.hub.load("zhiqwang/yolov5-rt-stack", "yolov5s", pretrained=True)

Loading checkpoint from official yolov5

The following is the interface for loading the checkpoint weights trained with ultralytics/yolov5. See our how-to-align-with-ultralytics-yolov5 notebook for more details.

from yolort.models import YOLOv5

# 'yolov5s.pt' is downloaded from https://github.com/ultralytics/yolov5/releases/download/v5.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)

:rocket: Deployment

Inference on LibTorch backend

We provide a notebook to demonstrate how the model is transformed into torchscript. And we provide an C++ example of how to infer with the transformed torchscript model. For details see the GitHub Actions.

Inference on ONNXRuntime backend

On the ONNXRuntime front you can use the C++ example, and we also provide a tutorial export-onnx-inference-onnxruntime for using the ONNXRuntime.

:art: Model Graph Visualization

Now, yolort can draw the model graph directly, checkout our model-graph-visualization notebook to see how to use and visualize the model graph.

YOLO model visualize

:mortar_board: Acknowledgement

  • The implementation of yolov5 borrow the code from ultralytics.
  • This repo borrows the architecture design and part of the code from torchvision.

:thinking: Contributing

See the CONTRIBUTING file for how to help out. BTW, leave a :star2: if you liked it, and this is the easiest way to support us :)

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

yolort-0.5.1.tar.gz (135.5 kB view hashes)

Uploaded Source

Built Distribution

yolort-0.5.1-py2.py3-none-any.whl (144.2 kB view hashes)

Uploaded Python 2 Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page