Export

Export the deploy-mode model + postprocessor as a single ONNX graph, plus a trtexec command builder for TensorRT.

dfine.export.onnx.export_onnx

export_onnx(model: Module, postprocessor: Module, file: str | Path, *, task: str = 'detect', imgsz: int = 640, batch: int = 1, opset: int = 16, dynamic: bool = True, simplify: bool = False, check: bool = True, device: device | str = 'cpu') -> Path

Export model (+ postprocessor) to an ONNX graph at file.

Parameters:

Name	Type	Description	Default
task	str	"detect" / "segment" / "sem_seg" — picks the output contract (see the module docstring). postprocessor is unused for "sem_seg".	'detect'
imgsz	int	square input resolution of the dummy input (match the model's imgsz).	640
batch	int	dummy batch size (a real value even when dynamic — used for tracing).	1
opset	int	ONNX opset (upstream uses 16).	16
dynamic	bool	mark the batch dim N dynamic on inputs and outputs.	True
simplify	bool	run onnxsim on the graph (needs onnxsim).	False
check	bool	run onnx.checker on the exported graph.	True
device	device | str	device to trace on.	'cpu'

Returns:

Type	Description
Path	The written .onnx :class:~pathlib.Path.

Source code in dfine/export/onnx.py

def export_onnx(
    model: nn.Module,
    postprocessor: nn.Module,
    file: str | Path,
    *,
    task: str = "detect",
    imgsz: int = 640,
    batch: int = 1,
    opset: int = 16,
    dynamic: bool = True,
    simplify: bool = False,
    check: bool = True,
    device: torch.device | str = "cpu",
) -> Path:
    """Export ``model`` (+ ``postprocessor``) to an ONNX graph at ``file``.

    Args:
        task: ``"detect"`` / ``"segment"`` / ``"sem_seg"`` — picks the output contract
            (see the module docstring). ``postprocessor`` is unused for ``"sem_seg"``.
        imgsz: square input resolution of the dummy input (match the model's ``imgsz``).
        batch: dummy batch size (a real value even when ``dynamic`` — used for tracing).
        opset: ONNX opset (upstream uses 16).
        dynamic: mark the batch dim ``N`` dynamic on inputs and outputs.
        simplify: run ``onnxsim`` on the graph (needs ``onnxsim``).
        check: run ``onnx.checker`` on the exported graph.
        device: device to trace on.

    Returns:
        The written ``.onnx`` :class:`~pathlib.Path`.
    """
    device = torch.device(device)
    file = Path(file)
    deploy, input_names, output_names = _build_deploy(model, postprocessor, task)
    deploy = deploy.to(device).eval()

    trace_batch = max(batch, 2) if dynamic else batch
    images = torch.rand(trace_batch, 3, imgsz, imgsz, device=device)
    sizes = torch.tensor([[imgsz, imgsz]] * trace_batch, device=device)
    args = (images, sizes) if "orig_target_sizes" in input_names else (images,)
    deploy(*args)

    dynamic_axes = None
    if dynamic:
        dynamic_axes = {name: {0: "N"} for name in input_names + output_names}

    export_kwargs = dict(
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=opset,
        do_constant_folding=True,
        verbose=False,
    )
    import inspect

    if "dynamo" in inspect.signature(torch.onnx.export).parameters:
        export_kwargs["dynamo"] = False

    file.parent.mkdir(parents=True, exist_ok=True)
    torch.onnx.export(deploy, args, str(file), **export_kwargs)

    if check:
        import onnx

        onnx.checker.check_model(onnx.load(str(file)))

    if simplify:
        import onnx
        import onnxsim

        shapes = {"images": tuple(images.shape)}
        if "orig_target_sizes" in input_names:
            shapes["orig_target_sizes"] = tuple(sizes.shape)
        simplified, ok = onnxsim.simplify(str(file), test_input_shapes=shapes)
        if not ok:
            raise RuntimeError(f"onnxsim failed to validate the simplified graph for {file}.")
        onnx.save(simplified, str(file))

    return file

dfine.export.onnx.tensorrt_command

tensorrt_command(onnx_file: str | Path, *, task: str = 'detect', imgsz: int = 640, fp16: bool = True, engine: str | None = None, max_batch: int = 32) -> str

Return the trtexec command to build a TensorRT engine from onnx_file.

The graph's batch dim is dynamic (H/W are fixed to the export resolution), so TensorRT needs an optimization profile; this provides min/opt/max shapes at imgsz (pass the same value you exported with) with batch 1..max_batch. The sem_seg graph has a single images input; detect/segment also take orig_target_sizes (pass the matching task). Run the returned command where trtexec (and OpenVINO's ovc <onnx_file> for OpenVINO) is installed — those toolchains are not Python deps here.

Source code in dfine/export/onnx.py

def tensorrt_command(
    onnx_file: str | Path,
    *,
    task: str = "detect",
    imgsz: int = 640,
    fp16: bool = True,
    engine: str | None = None,
    max_batch: int = 32,
) -> str:
    """Return the ``trtexec`` command to build a TensorRT engine from ``onnx_file``.

    The graph's batch dim is dynamic (H/W are fixed to the export resolution), so
    TensorRT needs an optimization profile; this provides min/opt/max shapes at
    ``imgsz`` (pass the same value you exported with) with batch 1..``max_batch``. The
    ``sem_seg`` graph has a single ``images`` input; ``detect``/``segment`` also take
    ``orig_target_sizes`` (pass the matching ``task``). Run the returned command where
    ``trtexec`` (and OpenVINO's ``ovc <onnx_file>`` for OpenVINO) is installed — those
    toolchains are not Python deps here.
    """
    onnx_file = Path(onnx_file)
    engine = engine or str(onnx_file.with_suffix(".engine"))
    hw = f"3x{imgsz}x{imgsz}"
    sizes = "" if task == "sem_seg" else ",orig_target_sizes:{b}x2"
    parts = [
        "trtexec",
        f"--onnx={onnx_file}",
        f"--saveEngine={engine}",
        f"--minShapes=images:1x{hw}{sizes.format(b=1)}",
        f"--optShapes=images:1x{hw}{sizes.format(b=1)}",
        f"--maxShapes=images:{max_batch}x{hw}{sizes.format(b=max_batch)}",
    ]
    if fp16:
        parts.append("--fp16")
    return " ".join(parts)