The YOLOS model was proposed in You Only Look at One Sequence: Rethinking Transformer in Vision through Object Detection by Yuxin Fang, Bencheng Liao, Xinggang Wang, Jiemin Fang, Jiyang Qi, Rui Wu, Jianwei Niu, Wenyu Liu. YOLOS proposes to just leverage the plain Vision Transformer (ViT) for object detection, inspired by DETR. It turns out that a base-sized encoder-only Transformer can also achieve 42 AP on COCO, similar to DETR and much more complex frameworks such as Faster R-CNN.

This guide shows how to export a trained YoloS model to an ONNX conforming to the requirements mentioned before.


Make sure to install the required packages:

pip install -r requirements.txt

Exporting the model to ONNX

  1. The following command exports a trained YoloS model to ONNX and saves it locally:

bash export-to-onnx.sh hustvl/yolos-tiny yolos-tiny
# This script is used to convert HuggingFace models to ONNX format opset version 12
set -e

# Check if the user has provided the model name
if [ -z "$2" ]
    echo "Usage: $0 <model_name> <folder_path>"
    exit 1

model_name=$1 # exampples: hustvl/yolos-tiny
folder_path=$2 # example: yolos-model

optimum-cli export onnx --model "$model_name" \
  --opset 12 \
  --framework pt \
  --batch_size 1 \
  --atol 0.001 \
  1. The exported ONNX needs a little bit of fiddling to make usable with the Nx AI Manager.

python complete_onnx.py
from os.path import dirname, abspath, join, splitext
from glob import glob
from utils import add_post_processing_to_onnx, simplify_onnx, set_input_shape, read_configs, update_onnx_doc_string

HERE = dirname(abspath(__file__))

if __name__ == '__main__':
    model_folder_path = join(HERE, 'yolos-tiny')

    onnx_path = glob(join(model_folder_path, '*.onnx'))[0]
    output_onnx_path = splitext(onnx_path)[0] + '-complete.onnx'

    metadata = read_configs(model_folder_path)
    image_size = metadata['image_size']
    if image_size is None:
        image_size = [416, 416]

    # Update input shape to static shape
    set_input_shape(onnx_path, output_onnx_path, (1, 3, image_size[0], image_size[1]))

    # Update ONNX doc string
    update_onnx_doc_string(output_onnx_path, metadata['means'], metadata['stds'])

    # Simplify  ONNX model
    simplify_onnx(output_onnx_path, output_onnx_path)

    add_post_processing_to_onnx(output_onnx_path, output_onnx_path, id2labels=metadata['id2label'],

    # Simplify ONNX model again
    simplify_onnx(output_onnx_path, output_onnx_path)
import json
import re

import numpy as np
import onnx
import sclblonnx as so
from onnx import save
from onnxsim import simplify

def add_post_processing_to_onnx(onnx_path: str, output_onnx_path: str, id2labels: dict[str, str],
                                image_size: tuple[int, int]):
    base_graph = so.graph_from_file(onnx_path)
    logits_name = base_graph.output[0].name
    boxes_name = base_graph.output[1].name
    input_name = base_graph.input[0].name

    transformation_matrix = np.array([[[1, 0, 1, 0], [0, 1, 0, 1], [-1/2, 0, 1/2, 0], [0, -1/2, 0, 1/2]]],
    so.add_constant(base_graph, 'transformation_matrix', transformation_matrix, 'FLOAT')

    class_ids = np.array(sorted(list(id2labels.keys())), dtype=np.int64)
    so.add_constant(base_graph, 'class_ids', class_ids, 'INT64')

    so.add_constant(base_graph, 'C1', np.array(1, dtype=np.int64), 'INT64')

    # Softmax the logits
    softmax_output = so.node('Softmax', inputs=[logits_name], outputs=['softmax_output'], axis=-1)  # [1, M, C]

    # Compute classes
    classes = so.node('ArgMax', inputs=['softmax_output'], outputs=['classes'], axis=-1, keepdims=1)  # [1, M, 1]
    float_classes = so.node('Cast', inputs=['classes'], outputs=['float_classes'], to=1)  # [1, M, 1]
    equal_classes = so.node('Equal', inputs=['classes', 'class_ids'], outputs=['equal_classes'])  # [1, M, C]
    equal_classes_int = so.node('Cast', inputs=['equal_classes'], outputs=['equal_classes_int'], to=7)  # [1, M, C]
    reducesum_classes = so.node('ReduceSum', inputs=['equal_classes_int'], outputs=['reducesum_classes'],
                                axes=(-1,), keepdims=1)  # [1, M, 1]
    mask_classes = so.node('Cast', inputs=['reducesum_classes'], outputs=['mask_classes'], to=9)  # [1, M, 1]

    # Compute scores
    scores = so.node('ReduceMax', inputs=['softmax_output'], outputs=['scores'], axes=(-1,), keepdims=1)  # [1, M, 1]
    mask_scores = so.node('Greater', inputs=['scores', 'nms_sensitivity-'], outputs=['mask_scores'])  # [1, M, 1]
    boxes_to_keep = so.node('And', inputs=['mask_classes', 'mask_scores'], outputs=['boxes_to_keep'])  # [1, M, 1]

    # box ids to keep
    box_ids = so.node('NonZero', inputs=['boxes_to_keep'], outputs=['box_ids'])  # [3, M]
    box_ids_row1 = so.node('Gather', inputs=['box_ids', 'C1'], outputs=['box_ids_row1'], axis=0)  # [M]

    # xywh to xyxy
    xyxy = so.node('MatMul', inputs=[boxes_name, 'transformation_matrix'],
                   outputs=['xyxy'])  # [1, M, 4]

    # Concat the boxes, score & classes
    bboxes = so.node('Concat', inputs=['xyxy', 'scores', 'float_classes'],
                     outputs=['bboxes'], axis=2)  # [1, M, 6]

    # Keep only the boxes to keep
    bboxes_to_keep = so.node('Gather', inputs=['bboxes', 'box_ids_row1'],
                             outputs=['bboxes_to_keep'], axis=1)  # [1, M, 6]
    bboxes_squeezed = so.node('Squeeze', inputs=['bboxes_to_keep'], outputs=['bboxes_squeezed'],
                              axes=(0,))  # [M, 6]

                 [softmax_output, classes, float_classes, equal_classes, equal_classes_int, reducesum_classes,
                  scores, mask_scores, boxes_to_keep, box_ids, box_ids_row1, xyxy, bboxes, bboxes_to_keep,

    # Add mask to the model
    so.delete_output(base_graph, boxes_name)
    so.delete_output(base_graph, logits_name)

    mask_bboxes(base_graph, 'bboxes_squeezed', 'mask-', image_size[1], image_size[0])
    so.add_output(base_graph, 'unmasked_bboxes', 'FLOAT', dimensions=[20, 6])

    so.add_input(base_graph, name='nms_sensitivity-', dimensions=[1], data_type='FLOAT')

    # Save the model
    so.graph_to_file(base_graph, output_onnx_path, onnx_opset_version=get_onnx_opset_version(onnx_path))

    # Rename model IO
    classes_str = ';'.join([f'{k}:{v}' for k, v in id2labels.items()])
    rename_io(output_onnx_path, output_onnx_path, **{input_name: 'image-',
                                                     'unmasked_bboxes': f'bboxes-format:xyxysc;{classes_str}'

def mask_bboxes(graph, bboxes_name, mask_name, w, h):
    so.add_input(graph, name=mask_name, dimensions=[h, w], data_type='BOOL')

    so.add_constant(graph, 'index_one_three', np.array([0, 2]), 'INT64')
    so.add_constant(graph, 'index_four', np.array([3, 3]), 'INT64')
    so.add_constant(graph, 'hw_clip_min', np.array(0), 'FLOAT')
    so.add_constant(graph, 'w_clip_max', np.array(w - 1), 'FLOAT')
    so.add_constant(graph, 'h_clip_max', np.array(h - 1), 'FLOAT')

    x_coordinates = so.node('Gather', inputs=[bboxes_name, 'index_one_three'], outputs=['x_coordinates'], axis=1)
    y_coordinates = so.node('Gather', inputs=[bboxes_name, 'index_four'], outputs=['y_coordinates'], axis=1)
    x_reducemean = so.node('ReduceMean', inputs=['x_coordinates'], outputs=['x_reducemean'], axes=(1,), keepdims=1)
    y_coordinate = so.node('ReduceMean', inputs=['y_coordinates'], outputs=['y_coordinate'], axes=(1,), keepdims=1)
    x_clipped = so.node('Clip', inputs=['x_reducemean', 'hw_clip_min', 'w_clip_max'], outputs=['x_clipped'])
    y_clipped = so.node('Clip', inputs=['y_coordinate', 'hw_clip_min', 'h_clip_max'], outputs=['y_clipped'])

    bottom_center_corner = so.node('Concat', inputs=['y_clipped', 'x_clipped'], outputs=['bottom_center_corner'],
    bottom_center_corner_int = so.node('Cast', inputs=['bottom_center_corner'], outputs=['bottom_center_corner_int'],
    bboxes_mask1 = so.node('GatherND', inputs=[mask_name, 'bottom_center_corner_int'], outputs=['bboxes_mask1'])
    bboxes_indices1 = so.node('NonZero', inputs=['bboxes_mask1'], outputs=['bboxes_indices1'])
    bboxes_indices1_squeezed = so.node('Squeeze', inputs=['bboxes_indices1'], outputs=['bboxes_indices1_squeezed'],
    new_bboxes = so.node('Gather', inputs=[bboxes_name, 'bboxes_indices1_squeezed'], outputs=['unmasked_bboxes'],

    so.add_nodes(graph, [x_coordinates, y_coordinates, x_reducemean, y_coordinate, y_clipped, x_clipped,
                         bboxes_mask1, bboxes_indices1,
                         bboxes_indices1_squeezed, new_bboxes])
    return graph

def rename_io(model_path, new_model_path=None, **io_names):
    if new_model_path is None:
        new_model_path = model_path

    g = so.graph_from_file(model_path)

    def log(old: bool = True):
        s = 'Old' if old else 'New'
        assert so.list_inputs(g)
        assert so.list_outputs(g)

    if io_names == {}:

    inputs = [i.name for i in g.input]
    outputs = [i.name for i in g.output]

    for k, v in io_names.items():
        pattern = re.compile(k)
        renamed = False

        for i in inputs:
            if pattern.match(i):
                renamed = True
                so.rename_input(g, i, v)

        if not renamed:
            for o in outputs:
                if pattern.match(o):
                    renamed = True
                    so.rename_output(g, o, v)

        if not renamed:


    so.graph_to_file(g, new_model_path, onnx_opset_version=get_onnx_opset_version(model_path))

def get_onnx_opset_version(onnx_path):
    model = onnx.load(onnx_path)
    opset_version = model.opset_import[0].version if len(model.opset_import) > 0 else 0
    return opset_version

def simplify_onnx(onnx_path: str, output_onnx_path: str):
        model, check = simplify(onnx_path, check_n=1)
        assert check, 'Failed to simplify ONNX model'
    except Exception as e:
        import traceback
        raise Exception('Failed to simplify ONNX model')

    save(model, output_onnx_path)

def update_onnx_doc_string(onnx_path: str, model_means: list[float], model_stds: list[float]):
    # Update the ONNX description
    graph = so.graph_from_file(onnx_path)
    # Add the model means and standard deviations to the ONNX graph description,
    # because that's used by the toolchain to populate some settings.
    graph.doc_string = json.dumps({'means': model_means, 'vars': model_stds})
    so.graph_to_file(graph, onnx_path, onnx_opset_version=12)

def set_input_shape(input_onnx: str, output_onnx: str, new_shape: tuple[int, int, int, int]):
    graph = so.graph_from_file(input_onnx)
    input_shape = graph.input[0].type.tensor_type.shape.dim
    for i, d in enumerate(input_shape):
        d.dim_value = new_shape[i]

    so.graph_to_file(graph, output_onnx, onnx_opset_version=12)

def read_configs(model_folder_path: str):
    from os.path import join
    import json

    preproccessing_config_path = join(model_folder_path, 'preprocessor_config.json')
    config_path = join(model_folder_path, 'config.json')

    with open(preproccessing_config_path, 'r') as f:
        preprocessor_config = json.load(f)

    with open(config_path, 'r') as f:
        config = json.load(f)

    metadata = {'means': None, 'stds': None, 'id2label': None, 'image_size': None}

    # Get model means and stds
    reescale_factor = preprocessor_config.get('rescale_factor', 1)
    do_rescale = preprocessor_config.get('do_rescale', False)
    image_mean = preprocessor_config.get('image_mean', [0, 0, 0])
    image_std = preprocessor_config.get('image_std', [1, 1, 1])
    if do_rescale:
        metadata['means'] = [round(m / reescale_factor, 2) for m in image_mean]
        metadata['stds'] = [round(s / reescale_factor, 2) for s in image_std]

    # Get model image size
    metadata['image_size'] = config.get('image_size', None)

    # Get model id2label
    id2label = config.get('id2label', None)
    id2label = {k: v for k, v in id2label.items() if v != 'N/A'}  # Eliminate N/A labels
    metadata['id2label'] = id2label

    return metadata
  1. Finally, here's a Python script that can be used in order to test the ONNX file using ONNXRuntime.

python test_onnx.py
from os.path import join, dirname, abspath

import cv2
import numpy as np
import onnxruntime as rt

PATH = dirname(abspath(__file__))

def test_model(model_path, img_path):
    sess = rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])

    # get input name
    input_name1 = sess.get_inputs()[0].name
    input_name2 = sess.get_inputs()[1].name
    input_name3 = sess.get_inputs()[2].name

    # get input dimensions
    input_shape = sess.get_inputs()[0].shape
    if input_shape[1] <= 3:  # nchw
        height, width = input_shape[2], input_shape[3]
    else:  # nhwc
        height, width = input_shape[1], input_shape[2]

    img = cv2.imread(img_path)
    img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # nwhc nchw
    img = (img.astype('float32') - np.array([123.67, 116.28, 103.53])) / np.array([58.395, 57.12, 57.375])
    img = np.transpose(img, (2, 0, 1)).astype('float32')
    img = np.expand_dims(img, axis=0)

    mask_area = np.repeat(1, width * height).astype('bool')
    mask_area = mask_area.reshape((height, width))
    mask_area[:, :width // 2] = 1  # mask the left part of the image

    bboxes = sess.run(None, {
        input_name1: img,
        input_name2: mask_area,
        input_name3: np.array([0.9]).astype('float32')
    bboxes = bboxes[-1]

    return bboxes, (width, height)

def visualize_bboxes(bboxes, img_path, width, height):
    img = cv2.imread(img_path)
    img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
    if bboxes[:, :4].max() <= 1.5:
        bboxes[:, :4] = bboxes[:, :4] * [width, height, width, height]
    for bbox in bboxes:
        x1, y1, x2, y2, score, class_id = bbox
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(img, f'{int(class_id)}', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    cv2.imshow('img', img)

if __name__ == '__main__':
    from glob import glob

    model_path = glob(join(PATH, '**', '*-complete.onnx'), recursive=True)[0]
    img_path = join(PATH, 'pedestrians.jpg')
    bboxes, (width, height) = test_model(model_path, img_path)
    visualize_bboxes(bboxes, img_path, width, height)

Beyond this example

This example is a starting point for exporting YoloS-tiny model to ONNX. Yet, this approach is valid for any YoloS model trained on any object detection dataset.

To adapt this example to your own model, you need to:

  • update the content of the "export-to-onnx.sh" and "complete_onnx.py" scripts accordingly.

Last updated