Package: yolort
YOLOv5 Runtime StackDocumentation•Installation Instructions•Deployment•Contributing•Reporting Issues🤗 IntroductionWhat it is.Yet another implementation of Ultralytics'sYOLOv5. yolort aims to make the training and inference of the object detection task 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 deploy the object detection more easily and friendly onLibTorch,ONNX Runtime,TVM,TensorRTand so on.About the code.Follow the design principle ofdetr:object detection should not be more difficult than classification, and should not require complex libraries for training and inference.yolortis very simple to implement and experiment with. Do you like the implementation of torchvision's faster-rcnn, retinanet or detr? Do you like yolov5? You'll loveyolort!🆕 What's NewDec. 27, 2021. AddTensorRTC++ interface example. Thanks toShiquan.Dec. 25, 2021. Support exporting toTensorRT, and inferencing withTensorRTPython interface.Sep. 24, 2021. AddONNX RuntimeC++ interface example. Thanks toFidan.Feb. 5, 2021. AddTVMcompile and inference notebooks.Nov. 21, 2020. Add graph visualization tools.Nov. 17, 2020. Support exporting toONNX, and inferencing withONNX RuntimePython interface.Nov. 16, 2020. Refactor YOLO modules and supportdynamic shape/batchinference.Nov. 4, 2020. AddLibTorchC++ inference example.Oct. 8, 2020. Support exporting toTorchScriptmodel.🛠️ UsageThere are no extra compiled components inyolortand package dependencies are minimal, so the code is very simple to use.Installation and Inference ExamplesAbove all, follow theofficial instructionsto install PyTorch 1.8.0+ and torchvision 0.9.0+Installation via pipSimple installation fromPyPIpipinstall-UyolortOr from Source# clone yolort repository locallygitclonehttps://github.com/zhiqwang/yolov5-rt-stack.gitcdyolov5-rt-stack# install in editable modepipinstall-e.Install pycocotools (for evaluation on COCO):pipinstall-U'git+https://github.com/ppwwyyxx/cocoapi.git#subdirectory=PythonAPI'To read a source of image(s) and detect its objects 🔥fromyolort.modelsimportyolov5s# Load modelmodel=yolov5s(pretrained=True,score_thresh=0.45)model.eval()# Perform inference on an image filepredictions=model.predict("bus.jpg")# Perform inference on a list of image filespredictions=model.predict(["bus.jpg","zidane.jpg"])Loading viatorch.hubThe models are also available via torch hub, to loadyolov5swith pretrained weights simply do:model=torch.hub.load("zhiqwang/yolov5-rt-stack:main","yolov5s",pretrained=True)Loading checkpoint from official yolov5The following is the interface for loading the checkpoint weights trained withultralytics/yolov5. Please see our documents on what weshareand how wedifferfrom yolov5 for more details.fromyolort.modelsimportYOLOv5# Download checkpoint from https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.ptckpt_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)🚀 DeploymentInference on LibTorch backendWe provide atutorialto demonstrate how the model is converted intotorchscript. And we provide aC++ exampleof how to do inference with the serializedtorchscriptmodel.Inference on ONNX Runtime backendWe provide a pipeline for deploying yolort with ONNX Runtime.fromyolort.runtimeimportPredictorORT# Load the serialized ONNX modelengine_path="yolov5n6.onnx"y_runtime=PredictorORT(engine_path,device="cpu")# Perform inference on an image filepredictions=y_runtime.predict("bus.jpg")Please check out thistutorialto use yolort's ONNX model conversion and ONNX Runtime inferencing. And you can use theexamplefor ONNX Runtime C++ interface.Inference on TensorRT backendThe pipeline for TensorRT deployment is also very easy to use.importtorchfromyolort.runtimeimportPredictorTRT# Load the serialized TensorRT engineengine_path="yolov5n6.engine"device=torch.device("cuda")y_runtime=PredictorTRT(engine_path,device=device)# Perform inference on an image filepredictions=y_runtime.predict("bus.jpg")Besides, we provide atutorialdetailing yolort's model conversion to TensorRT and the use of the Python interface. Please check thisexampleif you want to use the C++ interface.🎨 Model Graph VisualizationNow,yolortcan draw the model graph directly, checkout ourtutorialto see how to use and visualize the model graph.👋 ContributingWe love your input! Please see ourContributing Guideto get started and for how to help out. Thank you to all our contributors! If you like this project please consider ⭐ this repo, as it is the simplest way to support us.📖 Citing yolortIf you use yolort in your publication, please cite it by using the following BibTeX entry.@Misc{yolort2021,author={Zhiqiang Wang and Shiquan Yu and Fidan Kharrasov},title={yolort: A runtime stack for object detection on specialized accelerators},howpublished={\url{https://github.com/zhiqwang/yolov5-rt-stack}},year={2021}}🎓 AcknowledgementThe implementation ofyolov5borrow the code fromultralytics.This repo borrows the architecture design and part of the code fromtorchvision.