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, *, 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
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,
    *,
    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:
        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 = DeployModel(model, postprocessor).to(device).eval()

    # The decoder takes a batched code path only when batch > 1; trace with >= 2 for a
    # dynamic export so the graph generalizes to any N (it still serves batch 1). Upstream
    # traces with 32 for the same reason.
    trace_batch = max(batch, 2) if dynamic else batch
    images = torch.rand(trace_batch, 3, imgsz, imgsz, device=device)
    orig_target_sizes = torch.tensor([[imgsz, imgsz]] * trace_batch, device=device)
    deploy(images, orig_target_sizes)  # sanity forward before export

    dynamic_axes = None
    if dynamic:
        dynamic_axes = {name: {0: "N"} for name in ("images", "orig_target_sizes")}
        dynamic_axes.update({name: {0: "N"} for name in ("labels", "boxes", "scores")})

    # Force the legacy TorchScript exporter (what upstream validated at opset 16); newer
    # torch defaults to the dynamo exporter, which needs the extra ``onnxscript`` dep.
    export_kwargs = dict(
        input_names=["images", "orig_target_sizes"],
        output_names=["labels", "boxes", "scores"],
        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, (images, orig_target_sizes), 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),
            "orig_target_sizes": tuple(orig_target_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, *, 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. 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,
    *,
    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``. 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}"
    parts = [
        "trtexec",
        f"--onnx={onnx_file}",
        f"--saveEngine={engine}",
        f"--minShapes=images:1x{hw},orig_target_sizes:1x2",
        f"--optShapes=images:1x{hw},orig_target_sizes:1x2",
        f"--maxShapes=images:{max_batch}x{hw},orig_target_sizes:{max_batch}x2",
    ]
    if fp16:
        parts.append("--fp16")
    return " ".join(parts)