Is this page helpful?

ONNX Conversion Guide#

TensorRT-RTX uses the Open Neural Network Exchange (ONNX) format as its primary model input. Before you can build a TensorRT-RTX engine, your model must be exported to an .onnx file. This guide covers how to export models from common training frameworks and how to build ONNX models programmatically.

Note

TensorRT-RTX supports ONNX opsets 9–22 (inclusive). Not all operators and precisions within these opsets are supported. For a complete list of supported operators, refer to the Operator Support reference.

Exporting from a Training Framework#

Most users train models in PyTorch, TensorFlow, or Hugging Face and export to ONNX. Choose the method that matches your framework.

Use torch.onnx.export() to convert a PyTorch model to ONNX:

python -m pip install torch onnx onnxscript

import onnx
import torch

class TinyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = torch.nn.Conv2d(
            3, 8, kernel_size=3, padding=1)

    def forward(self, x):
        return self.conv(x)

model = TinyModel().eval()
dummy_input = torch.randn(1, 3, 224, 224)
batch_size = torch.export.Dim("batch_size")

torch.onnx.export(model, dummy_input, "model.onnx",
                  input_names=["input"],
                  output_names=["output"],
                  dynamic_shapes={"x": {0: batch_size}})

onnx.checker.check_model(onnx.load("model.onnx"))

For detailed options and troubleshooting, refer to the PyTorch ONNX export documentation and tutorial.

TensorFlow does not include built-in ONNX support. Use the open-source tf2onnx tool:

pip install tf2onnx
python -m tf2onnx.convert --saved-model ./saved_model --output model.onnx

Use the Optimum library to export Hugging Face model checkpoints to ONNX:

pip install optimum[onnxruntime]
optimum-cli export onnx --model bert-base-uncased ./bert_onnx/

Building ONNX Models Programmatically#

If your framework does not support ONNX export, you can construct ONNX models directly using the ONNX Python API. This approach defines the model graph, operators, and weights using protocol buffers.

Example: Simple multilayer perceptron
import onnx
import onnx.helper as helper
import onnx.numpy_helper as numpy_helper
from onnx import TensorProto
nb_inputs, nb_hidden, nb_outputs = 10, 20, 1

# Create input and output tensors
input = helper.make_tensor_value_info('input', TensorProto.FLOAT, ['batch', nb_inputs])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, ['batch', nb_outputs])

# Create weights and biases as numpy tensors (use trained values in production)
import numpy as np
W1 = np.random.randn(nb_hidden, nb_inputs).astype(np.float32)
B1 = np.random.randn(nb_hidden).astype(np.float32)
W2 = np.random.randn(nb_outputs, nb_hidden).astype(np.float32)
B2 = np.random.randn(nb_outputs).astype(np.float32)
init_W1 = numpy_helper.from_array(W1, name='W1')
init_B1 = numpy_helper.from_array(B1, name='B1')
init_W2 = numpy_helper.from_array(W2, name='W2')
init_B2 = numpy_helper.from_array(B2, name='B2')

# Define operator nodes
gemm1 = helper.make_node('Gemm', ['input','W1','B1'], ['hidden1'], alpha=1.0, beta=1.0, transB=1)
relu1 = helper.make_node('Relu', ['hidden1'], ['relu1'])
gemm2 = helper.make_node('Gemm', ['relu1','W2','B2'], ['output'], alpha=1.0, beta=1.0, transB=1)

# Assemble graph
graph = helper.make_graph(nodes=[gemm1, relu1, gemm2],
                        name='MLP',
                        inputs=[input],
                        outputs=[output],
                        initializer=[init_W1, init_B1, init_W2, init_B2])
model = helper.make_model(graph, producer_name='simple_mlp')

# Validate and save
onnx.checker.check_model(model)
onnx.save(model, 'simple_mlp.onnx')

Dynamic Shape Considerations#

A model exported with dynamic shapes can contain operators whose output rank depends on dimensions that are not known until runtime. TensorRT-RTX resolves the output rank when it builds the engine, so these operators carry additional rules.

Squeeze with omitted axes

TensorRT-RTX supports ONNX Squeeze with omitted axes on inputs with dynamic shapes. The parser removes the dimensions that are known to be 1 in the network definition, and the dimensions that remain must not become 1 at runtime.

Releases before TensorRT-RTX 1.7 rejected this pattern with the following parser error.

Cannot infer squeeze dimensions from a dynamic shape!

Not every input is supported. If you still encounter this error, set the axes input of the Squeeze node explicitly when you export the model, so that the parser does not need to infer which dimensions to remove.

For the other ONNX parser fixes in this release, refer to the TensorRT-RTX Release Notes.

Next Steps#

After you have an .onnx file, proceed to build and run a TensorRT-RTX engine: