CloudSoftwareHardware
Software Stack
DepthAI
OAK Apps
AI Models

ON THIS PAGE

  • Conversion to ONNX
  • Overview
  • Conversion
  • From PyTorch
  • YOLO Models
  • From TensorFlow
  • Verification
  • Input/Output Tensors
  • Performance

Conversion to ONNX

Overview

Open Neural Network Exchange (ONNX) is a widely used format to represent machine learning models. It's open-source and supported by a wide range of frameworks and tools. Here we collect some of the most common approaches of converting your model to .onnx. It's adviced to do so as it opens up the most options at conversion for the appropriate RVC Platform.

Conversion

From PyTorch

You can utilize the PyTorch ONNX API to convert and export your model:
Python
1import torch
2import torch.nn as nn
3
4# Define the model architecture (must match the saved model)
5class Model(nn.Module):
6    ...
7    
8# Initialize the model and load the trained weights
9model = Model()
10model.load_state_dict(torch.load("model_name.pt"))
11model.eval()  # Set the model to inference mode
12
13# Define the input shape and create a dummy input tensor
14input_shape = ... # e.g. (1, 3, 512, 288)
15dummy_input = torch.randn(input_shape)
16
17# Export the model to ONNX format
18torch.onnx.export(model, dummy_input, "model_name.onnx")
This type of conversion does not alter the model’s architecture; it only changes the format in which the model is saved. As a result, the converted model exhibits performance identical to that of the original. For more information, refer to the official PyTorch export tutorial.

YOLO Models

For the conversion of YOLO models, we recommend using our tools-cli package—a specialized utility that streamlines the conversion process. It supports converting YOLO models from versions V5 through V11, including Gold YOLO, and covers advanced variants such as oriented bounding box (OBB) detection, pose estimation, and instance segmentation for YOLOv8 and YOLO11. Unlike the general PyTorch conversion, this process modifies the model architecture to standardize its outputs. This allows for out-of-the-box parsing using the native DepthAI DetectionNetwork node. (or the YOLOExtendedParser parser node, for pose and instnace segmentation models). As a result, post-processing can be performed entirely on-device, eliminating host-device data transfers and reducing overall latency. While this is the recommended approach, you can still export YOLO models using the general conversion method. However, keep in mind that you will need to handle output parsing manually, as native support is not provided. Additionally, be aware that the exported ONNX model may not be directly convertible or runnable due to limitations in supported operations—both in the converter and on the device.You can install the tools-cli package, as:
Command Line
1git clone --recursive https://github.com/luxonis/tools.git
2cd tools
3pip install .
Once installed, you can run the conversion directly from the command line (just make sure you are runnig from the root of the tools directory):
Command Line
1tools <MODEL>.pt --imgsz "<WIDTH> <HEIGHT>"
You can find more information about conversion parameters in the tools-cli documentation. For example, to convert a YOLOv6n model that expects BGR input images with a resolution of (512, 288), run the following:
Command Line
1tools yolov6n.pt --imgsz "512 288" --encoding BGR
The resulting ONNX model has input shape of (1, 3, 512, 288), and the following output shapes: (1,85,36,64), (1,85,18,32), and (1,85,9,16) (in contrast to the original model’s output shapes as obtained by the general conversion: (1, 2304, 85), (1, 576, 85), and (1, 144, 85); or (1, 3024, 85) if concatenated into a single tensor).

From TensorFlow

For models in in TensorFlow (.pb), Keras (.h5), tensorflow.js (.json and .bin) or TensorFlow Lite (.tflite) format we recommend using the tf2onnx conversion tool.
  • First, install the tensorflow and tf2onnx packages:
Command Line
1pip install tensorflow
2pip install -U tf2onnx
  • Second, the conversion can be done via the command line:
Command Line
1python -m tf2onnx.convert --saved-model tensorflow-model-path --output model.onnx

Verification

Input/Output Tensors

Prior converting ONNX for RVC Platform, make sure to check that input/output tensors an ONNX model are:
  • Correctly defined (taking inputs from / sending outputs to the appropriate layers).
    • If not, use the onnx-modifier tool to re-define input/output tensors.
  • Shape is of the form NCHW (batch-size, color-channels, height, width).
    • If not, define --inputs-as-nchw data and --outputs-as-nchw flags when running the tf2onnx tool.
  • None of the shape dimensions is dynamic (e.g. that batch size is fixed).
    • If dynamic, make shapes fixed using the onnxruntime library (see this tutorial for more information).

Performance

Converting a model to ONNX can sometimes result in slight differences in model performance. These discrepancies are generally due to differences in rounding, numerical precision, layer implementations, and operation approximations. It's adviced to compare the outputs of the original and ONNX models using various test inputs. You can evaluate the differences using metrics like mean absolute error (MAE) or mean squared error (MSE).If you identify discrepancies, consider the following steps:
  • Update to the Latest ONNX Opset: Ensure you are using the latest ONNX opset version. Newer opsets often include improvements and bug fixes that enhance model compatibility.
  • Simplify Model Operations: Simplify complex operations or layers in your model before converting it to ONNX. This can help reduce the likelihood of conversion-related issues.
  • Ensure Precision Consistency: Use the same data types (e.g., float32) for weights, inputs, and outputs in both models to maintain numerical precision.
By following these steps, you can minimize discrepancies and ensure that your ONNX model closely matches the performance of your original model.