Skip to content

Exporting to ONNX

echelon3 export turns a trained checkpoint into a single self-contained ONNX graph. The key idea: the exporter fuses your preprocess, the network, and an optional postprocess into one model, so the exported graph accepts a raw image tensor and returns final outputs — no Python preprocessing needed at inference time.

Install the export extra

ONNX export needs onnx and onnxruntime:

pip install echelon3[export]

The export section

create_exporters reads three things from the export section:

  • preprocess — an ordered map of torch.nn.Modules (identical shape to transform.*.preprocess). Defaults to Identity if omitted.
  • postprocess — the same, applied to the network's output. Defaults to Identity.
  • exporters — a map of named exporter triples. Each one is built with the network plus the fused preprocess/postprocess and then run.

The exporter base class wraps everything into one module whose forward pass is:

preprocess  →  net  →  postprocess

That fused module is what gets written to ONNX. Put pixel scaling / normalization in preprocess and any decode step (softmax-to-score, argmax-to-mask, sigmoid-threshold) in postprocess, and the graph becomes drop-in for a runtime that feeds it raw images.

OnnxExporter keys

module: echelon3.exporters.onnx / type: OnnxExporter:

Key Meaning
target Output .onnx path (parent directories are created).
input_names Names for the graph inputs.
output_names Names for the graph outputs.
input_shape Shape of the example input used to trace/script (e.g. [1, 3, 64, 64]).
opset ONNX opset version (default 18).
use_tracing true traces the module; false runs torch.jit.script first (default false).
dynamic_axes Optional map marking axes as dynamic (e.g. a variable batch size).
do_constant_folding Fold constants during export (default true).
use_aten_fallback Export with the ATen fallback operator set (default false).

The example input is uint8

The exporter builds its dummy input as a uint8 tensor in [0, 255] with input_shape. That is why preprocessing (e.g. To01) belongs inside the exported graph: the resulting ONNX model consumes a raw uint8 NCHW image directly.

Running the export

echelon3 export --config-dir configs --config-name my_experiment

The exporter builds the network from net, and — if a target section is present — loads the latest checkpoint from target.path (a legacy module. prefix from DDP checkpoints is stripped automatically). It then runs every exporter under export.exporters in turn. With no target, it exports the freshly-initialized network (useful for shape/graph checks).

Example

The export section from examples/configs/smoke.yaml folds a To01 step into the graph and writes one ONNX file:

export:
  preprocess:
    to01:
      name: to01
      module: echelon3.transforms.preprocess.basic
      type: To01
  exporters:
    onnx:
      module: echelon3.exporters.onnx
      type: OnnxExporter
      config:
        target: ${oc.env:SMOKE_TARGET,./targets/smoke}/smoke.onnx
        input_names: [images]
        output_names: [logits]
        input_shape: [1, 3, 64, 64]
        use_tracing: true
        opset: 18

Run it after training the smoke model:

SMOKE_TARGET=./targets/smoke \
    echelon3 export --config-dir examples/configs --config-name smoke

This writes ./targets/smoke/smoke.onnx. The graph takes a uint8 images input of shape [1, 3, 64, 64], scales it to [0, 1] internally via the fused To01, runs the classifier, and returns logits.

Dynamic batch size

To accept any batch size at inference, mark the batch axis dynamic:

config:
  input_shape: [1, 3, 64, 64]
  dynamic_axes:
    images: { 0: batch }
    logits: { 0: batch }

Adding a postprocess

To emit class scores instead of raw logits, append a postprocess step:

export:
  preprocess:
    to01: { name: to01, module: echelon3.transforms.preprocess.basic, type: To01 }
  postprocess:
    score: { name: score, module: echelon3.transforms.preprocess.basic, type: ScoreFromLogits, config: { score_class: 1 } }
  exporters:
    onnx: { ... }

For segmentation, echelon3.transforms.preprocess.segmentation.ToSegmentationMask or ToBinarySegmentationMask make good postprocess steps.

Next