Skip to main content

Yet Another YOLOv5 and its Additional Runtime Stack

Project description

🔦 yolort - YOLOv5 Runtime Stack

CI testing PyPI version codecov

What it is. Yet another implementation of Ultralytics's yolov5, and with modules refactoring to make it available in deployment backends such as 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

🆕 What's New

  • Support exporting to TorchScript model. Oct. 8, 2020.
  • Support inferring with LibTorch cpp interface. Oct. 10, 2020.
  • Add TorchScript cpp inference example. Nov. 4, 2020.
  • Refactor YOLO modules and support dynmaic batching inference. Nov. 16, 2020.
  • Support exporting to ONNX, and inferring with ONNXRuntime interface. Nov. 17, 2020.
  • Add graph visualization tools. Nov. 21, 2020.
  • Add TVM compile and inference notebooks. Feb. 5, 2021.

🛠️ 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

  • Installation via Pip

    Simple installation from PyPI

    pip install -U yolort
    

    Or from Source

    # clone flash repository locally
    git clone https://github.com/zhiqwang/yolov5-rt-stack.git
    cd yolov5-rt-stack
    # install in editable mode
    pip install -e .
    
  • To read a source of image(s) and detect its objects 🔥

    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)

Updating checkpoint from ultralytics/yolov5

The module state of yolort has some differences comparing to ultralytics/yolov5. We can load ultralytics's trained model checkpoint with minor changes, and we have converted ultralytics's lastest release v3.1 checkpoint here.

Expand to see more information of how to update ultralytics's trained (or your own) model checkpoint.
  • If you train your model using ultralytics's repo, you should update the model checkpoint first. ultralytics's trained model has a limitation that their model must load in the root path of ultralytics, so a important thing is to desensitize the path dependence as follows:

    # Noted that current path is the root of ultralytics/yolov5, and the checkpoint is
    # downloaded from <https://github.com/ultralytics/yolov5/releases/download/v3.1/yolov5s.pt>
    ultralytics_weights = 'https://github.com/ultralytics/yolov5/releases/download/v3.1/yolov5s.pt'
    checkpoints_ = torch.load(ultralytics_weights, map_location='cpu')['model']
    torch.save(checkpoints_.state_dict(), desensitize_ultralytics_weights)
    
  • Load yolort model as follows:

    from hubconf import yolov5s
    
    model = yolov5s()
    model.eval()
    
  • Now let's update ultralytics/yolov5 trained checkpoint, see the conversion script for more information:

    from utils.updated_checkpoint import update_ultralytics_checkpoints
    
    model = update_ultralytics_checkpoints(model, desensitize_ultralytics_weights)
    # updated checkpint is saved to checkpoint_path_rt_stack
    torch.save(model.state_dict(), checkpoint_path_rt_stack)
    

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.

🎨 Model Graph Visualization

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

YOLO model visualize

🎓 Acknowledgement

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

🤗 Contributing

We appreciate all contributions. If you are planning to contribute back bug-fixes, please do so without any further discussion. If you plan to contribute new features, utility functions or extensions, please first open an issue and discuss the feature with us. BTW, leave a 🌟 if you liked it, this means a lot to 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.3.1.tar.gz (69.3 kB view details)

Uploaded Source

Built Distribution

yolort-0.3.1-py2.py3-none-any.whl (71.2 kB view details)

Uploaded Python 2Python 3

File details

Details for the file yolort-0.3.1.tar.gz.

File metadata

  • Download URL: yolort-0.3.1.tar.gz
  • Upload date:
  • Size: 69.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.3.0 pkginfo/1.7.0 requests/2.25.1 setuptools/53.0.0 requests-toolbelt/0.9.1 tqdm/4.55.1 CPython/3.9.1

File hashes

Hashes for yolort-0.3.1.tar.gz
Algorithm Hash digest
SHA256 308135b8aae1bc9293fa620d957862ec8aa59c6a208b2e21aced1f56b7990d41
MD5 243f652ae2f2f49c00b3dd20740f7b4e
BLAKE2b-256 a5956d673ea25d29cfb8a89be09ed59072c7978dc1fd610a37d267e21ac1804d

See more details on using hashes here.

File details

Details for the file yolort-0.3.1-py2.py3-none-any.whl.

File metadata

  • Download URL: yolort-0.3.1-py2.py3-none-any.whl
  • Upload date:
  • Size: 71.2 kB
  • Tags: Python 2, Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.3.0 pkginfo/1.7.0 requests/2.25.1 setuptools/53.0.0 requests-toolbelt/0.9.1 tqdm/4.55.1 CPython/3.9.1

File hashes

Hashes for yolort-0.3.1-py2.py3-none-any.whl
Algorithm Hash digest
SHA256 0e0bb5d6ce0defe80ea340996c2918004ca51f40f17641db193c24af5d3a5f4c
MD5 e5c53b18265a6f9c8421412f85e679ef
BLAKE2b-256 7fb40143dfdcba8b21021771970a16b0a654ea2f5f7704816cd7ca1cbc59c5b8

See more details on using hashes here.

Supported by

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