Why yolort?

• TorchVision coding style, easy to understand, easy to adapt at will.

• Seamlessly load the official YOLOv5 trained checkpoints, ensuring consistency with the YOLOv5 structure.

• Full unit testing.

• Deployment-friendly, TensorRT / ONNX Runtime / LibTorch / ncnn multi-backend support.

Install yolort from PyPI

%pip install -q yolort

Note: you may need to restart the kernel to use updated packages.

from typing import List

import torch
from torch import nn, Tensor

Part of the code details

1. LetterBox is implemented using PyTorch native operators, via torch.nn.functional.interpolate and torch.nn.functional.pad operators.

2. Easily exported to TorchScript and ONNX formats.

3. P5/P6-compatible PathAggregationNetwork implementation, which fully captures the symmetry in the PAN structure.

4. Minimalist implementation of PostProcess.

Compatible with YOLOv5 original model structure

from yolort.v5 import attempt_download
from yolort.models._checkpoint import load_from_ultralytics

# yolov5s6.pt is downloaded from 'https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s6.pt'
model_path = "yolov5s6.pt"

checkpoint_path = attempt_download(model_path)
version = "r6.0"

model_info = load_from_ultralytics(checkpoint_path, version=version)

Obtain Backbone

from yolort.models.backbone_utils import darknet_pan_backbone

backbone_name = f"darknet_{model_info['size']}_{version.replace('.', '_')}"
depth_multiple = model_info["depth_multiple"]
width_multiple = model_info["width_multiple"]
use_p6 = model_info["use_p6"]
backbone = darknet_pan_backbone(backbone_name, depth_multiple, width_multiple, version=version, use_p6=use_p6)

Obtain full YOLOv5 model structure

from yolort.models.box_head import YOLOHead

class YOLO(nn.Module):
    def __init__(self, backbone: nn.Module, strides, num_anchors, num_classes: int):
        super().__init__()
        self.backbone = backbone
        self.head = YOLOHead(backbone.out_channels, num_anchors, strides, num_classes)

    def forward(self, samples: Tensor) -> List[Tensor]:

        # get the features from the backbone
        features = self.backbone(samples)

        # compute the yolo heads outputs using the features
        head_outputs = self.head(features)
        return head_outputs

strides = model_info["strides"]
num_anchors = len(model_info["anchor_grids"][0]) // 2
num_classes = model_info["num_classes"]
model = YOLO(backbone, strides, num_anchors, num_classes)

model.load_state_dict(model_info["state_dict"])

model = model.eval()

inputs = torch.rand(1, 3, 640, 640)

outs = model(inputs)

Functional deployment support solutions

fx_trace_model = torch.fx.symbolic_trace(model)

jit_trace_model = torch.jit.script(model)

View this document as a notebook: https://github.com/zhiqwang/yolort/blob/main/notebooks/why-yolort.ipynb

---