Conversion to ONNX
Overview
.onnx. It's adviced to do so as it opens up the most options at conversion for the appropriate RVC Platform.Conversion
From PyTorch
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")YOLO Models
tools-cli package, as:Command Line
1git clone --recursive https://github.com/luxonis/tools.git
2cd tools
3pip install .tools directory):Command Line
1tools <MODEL>.pt --imgsz "<WIDTH> <HEIGHT>"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(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).If you are interested in running the tools using Docker, please check out these guidelines.
From TensorFlow
.pb), Keras (.h5), tensorflow.js (.json and .bin) or TensorFlow Lite (.tflite) format we recommend using the tf2onnx conversion tool.- First, install the
tensorflowandtf2onnxpackages:
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.onnxVerification
Input/Output Tensors
- 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 dataand--outputs-as-nchwflags when running the tf2onnx tool.
- If not, define
- 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).
We suggest using the Netron tool to inspect the model.
Performance
- 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.