Export
Exports are written to runs/export/exp, then exp2, exp3, and so on by
default. Use project and name, an exact save_dir, or output for an exact
artifact filename. Set exist_ok=True only when replacing/reusing an explicit
destination is intentional. Export files are atomically published, and each run
includes args.yaml and environment.yaml.
Export a nitid checkpoint or supported model name to ONNX, OpenVINO IR, TorchScript, or TensorRT for deployment. Format availability depends on the task. To run the exported artifact outside nitid, see Deployment.
| Task | ONNX | OpenVINO | TorchScript | TensorRT |
|---|---|---|---|---|
| Detection | yes | yes | yes | yes |
| Instance segmentation | yes | yes | yes | yes |
| Semantic segmentation | yes | yes | no | no |
ONNX
from nitid import NITID
model = NITID("model1l", task="detect")
model.export(format="onnx")
# → dfine_640.onnx
uv run nitid export model=model1l task=detect format=onnx
The exported model takes a single input images [B, 3, H, W]. Detection returns
(labels, boxes, scores) with the postprocessor in deploy mode baked in.
Instance segmentation returns (labels, boxes, scores, masks); masks
contains selected mask probabilities aligned with the first three outputs.
Semantic segmentation exports return one output, semantic_logits [B, C, H,
W], at model-input resolution. Apply softmax for per-class probabilities and
argmax over axis 1 for the class-ID map. The integration test executes this
graph with ONNX Runtime and checks numerical parity with PyTorch.
semantic = NITID("semantic_best.pth", task="semantic")
semantic.export(format="onnx", output="semantic.onnx")
Semantic ONNX and OpenVINO exports are intended for dense-map deployment. Apply the same softmax/argmax post-processing after either backend.
Options
| Argument | Default | Description |
|---|---|---|
imgsz |
640 | Input spatial size. Use the checkpoint's trained/eval size unless you intentionally trained another size. |
batch |
1 | Static batch size (use dynamic=True for variable batch) |
dynamic |
False |
Export with dynamic batch and spatial axes |
simplify |
True |
Run onnxsim to fold constants and simplify the graph |
opset |
17 | ONNX opset version |
Dynamic export
model.export(format="onnx", dynamic=True, simplify=False)
Dynamic export currently marks the batch axis as dynamic.
OpenVINO
Install the optional OpenVINO dependency:
uv sync --extra openvino
Export through the corrected ONNX graph to OpenVINO Intermediate Representation (IR):
from nitid import NITID
model = NITID("model1l", task="detect")
xml_path = model.export(format="openvino")
# → runs/export/exp/dfine_640.xml
# → runs/export/exp/dfine_640.bin
uv run nitid export model=model1l task=detect format=openvino
The .xml file stores the graph and the matching .bin file stores its weights.
export() returns the .xml path. nitid first exports its deployment graph to a
temporary ONNX file, converts it with openvino.convert_model(), verifies that
OpenVINO can load and compile the IR on CPU, and then publishes the .bin and
.xml files. The temporary ONNX file is removed automatically.
Load the result for inference with OpenVINO Runtime:
import openvino as ov
core = ov.Core()
compiled_model = core.compile_model(xml_path, "AUTO")
outputs = compiled_model([images])
OpenVINO export uses the ONNX-related imgsz, batch, dynamic, simplify,
and opset arguments. Set half=True to compress IR weights to FP16; the
default preserves FP32 weights to make numerical comparison with ONNX easier.
Detection, instance segmentation, and semantic segmentation all route through
this ONNX-to-IR path.
TorchScript
model.export(format="torchscript")
# → dfine_640.torchscript
TorchScript export uses torch.jit.trace with strict=False because
D-FINE's forward pass returns a dict.
TensorRT
pip install --extra-index-url https://pypi.nvidia.com tensorrt>=8.6
model.export(format="tensorrt")
# → dfine_640.engine
uv run nitid export model=model1l task=detect format=tensorrt
The workflow is: trace model → temporary ONNX → TensorRT engine (the intermediate ONNX is removed automatically). Requires a CUDA-capable GPU.
Options
| Argument | Default | Description |
|---|---|---|
imgsz |
640 | Input spatial size |
batch |
1 | Static batch size (use dynamic=True for variable batch) |
dynamic |
False |
Build engine with dynamic batch axis (min=1, opt=batch, max=batch×4) |
half |
False |
Enable FP16 precision |
FP16
model.export(format="tensorrt", half=True)
FP16 halves memory usage and typically improves throughput on modern GPUs. nitid warns if the GPU reports no native FP16 support but still builds the engine.
Dynamic batch
model.export(format="tensorrt", dynamic=True, batch=4)
# min=1, opt=4, max=16
Raw export (no postprocessor)
By default the exported graph has the postprocessor baked in, so it returns the
decoded (labels, boxes, scores, ...) contract described above. Pass
postprocess=False to export the decoder's raw outputs instead and do the
decoding in the consuming application:
model.export(format="openvino", postprocess=False)
model.export(format="onnx", postprocess=False)
uv run nitid export model=model1l task=detect format=openvino postprocess=false
This is the same graph the backend="openvino" inference path compiles in
memory, only written to disk as a regular export artifact. Use it when the
consumer already has its own (often batched or fused) postprocessing, or when it
needs every decoder query rather than the top-k the postprocessor keeps.
Output names per task:
| Task | Raw outputs |
|---|---|
| Detection | pred_logits [B, Q, C], pred_boxes [B, Q, 4] |
| Instance segmentation | pred_logits, pred_boxes, pred_masks |
| Semantic segmentation | sem_seg_logits [B, C, H, W] |
Q is the number of decoder queries (300 for the official checkpoints) — all of
them, with no score threshold or top-k applied. pred_logits are unactivated,
and pred_boxes are cxcywh normalised to [0, 1].
The consumer has to reproduce what the postprocessor would have done:
scores = sigmoid(pred_logits)(focal-loss models; softmax otherwise)topkover the flattened[Q * C]scores, thenlabel = index % Candquery = index // C- convert the gathered boxes from
cxcywhtoxyxy - scale by the input
[w, h, w, h]
No NMS is involved — D-FINE is NMS-free. The reference implementation is
DFINEPostProcessor.forward() in
dfine/nn/postprocessor.py,
which also decodes the instance masks.
For semantic segmentation the raw export returns the same dense logit map, only
named sem_seg_logits instead of semantic_logits; apply softmax/argmax as
usual.
postprocess=False applies to ONNX, OpenVINO, and TensorRT. TorchScript always
traces the raw model (it has no postprocessor to drop), so the flag makes no
difference there.
Export spatial size
Most checkpoints are calibrated for a fixed evaluation spatial size
(eval_spatial_size, usually [640, 640]). Use the checkpoint's configured
size unless you have intentionally trained/evaluated another size.
print(model._cfg["eval_spatial_size"]) # [640, 640]
model.export(format="onnx", imgsz=640)