📦 Export¶
Writing a trained checkpoint as an ONNX file: graph_from_module rebuilds a task's suppression-free one-to-one graph from a loaded checkpoint, and export_graph writes it. The lucid-export command wraps both; how to run it is under Exporting to ONNX. Generated from the docstrings in src/lucid_yolo/export.py.
From a checkpoint¶
graph_from_module(module: DetectionLitModule, canvas: tuple[int, int], *, device: device | None = None) -> E2EExportGraph
Build a checkpoint-loaded module's export graph, its weights checked key for key.
A Lightning checkpoint loads into a :class:~lucid_yolo.ptl.module.DetectionLitModule
(:func:~lucid_yolo.eval.checkpoint.load_eval_module), and the export graphs wrap a composite model's
deploy() view, which the Lightning module does not have. This is the bridge between the two. It reads the
task from the module, rebuilds that task's composite — :class:~lucid_yolo.models.build.Detector,
:class:~lucid_yolo.models.build.Segmenter, :class:~lucid_yolo.models.build.OrientedDetector or
:class:~lucid_yolo.models.build.KeypointDetector — at the module's own multipliers, class count and point
count, copies the module's weights into it, and wraps its deploy() view in the task's graph class.
The rebuild is at raw multipliers rather than a variant name, so a module built outside the registry's five
rows exports as readily as one the training CLI configured. The copy is checked strictly in both directions:
the module must supply every key the composite holds and nothing else, apart from the training-only loss state
in :data:_TRAINING_ONLY_PREFIXES. That works because the module keeps its stages as flat
backbone/neck/head attributes (plus the three segmentation stages), the same names the composites
use, so the two state dicts share their keys.
The graph owns a copy of the weights: module is read and not modified, and stays on whatever device it was
on. Whichever weights module holds are what is exported, so a module
:func:~lucid_yolo.eval.checkpoint.load_eval_module returned with use_ema=True exports its EMA shadow.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
module
|
DetectionLitModule
|
The module to export, as
:func: |
required |
canvas
|
tuple[int, int]
|
Input canvas |
required |
device
|
device | None
|
Device the graph is built on, as :class: |
None
|
Returns:
| Type | Description |
|---|---|
E2EExportGraph
|
The task's eval-mode graph, ready for :func: |
Raises:
| Type | Description |
|---|---|
ValueError
|
If the module's task is not one of |
Examples:
>>> import torch
>>> from lucid_yolo.ptl.module import DetectionLitModule
>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=64, num_classes=2).eval()
>>> graph = graph_from_module(module, canvas=(64, 64))
>>> type(graph).__name__, graph.training
('DetectExportGraph', False)
>>> with torch.no_grad():
... detections = graph(torch.zeros(1, 3, 64, 64))
>>> detections.shape
torch.Size([1, 300, 6])
export_graph(graph: E2EExportGraph, path: str | Path, *, example: Tensor | None = None, opset_version: int = ONNX_OPSET) -> Path
Write one export graph to an ONNX file at a pinned opset.
The shipped counterpart to the four graph classes above: they compose the decode, this emits the file. It
exists so the opset is pinned somewhere a consumer runs rather than left to whatever
:func:torch.onnx.export's default is on the torch a given machine installed — see :data:ONNX_OPSET for what
that pin buys. The export gate calls this function rather than :func:torch.onnx.export directly, so the graph
an operator gets is the graph the gate certified.
The emitted file is static in every axis: no dynamic_axes is passed, so it answers for exactly the shape of
example — the canvas graph baked into its buffers, and that example's batch. A consumer needing a second
shape exports a second file.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
graph
|
E2EExportGraph
|
A built task graph — :class: |
required |
path
|
str | Path
|
Destination |
required |
example
|
Tensor | None
|
Input the graph is traced on, of shape |
None
|
opset_version
|
int
|
ONNX opset the graph is emitted at. Defaults to
:data: |
ONNX_OPSET
|
Returns:
| Type | Description |
|---|---|
Path
|
The written path. |
Raises:
| Type | Description |
|---|---|
OnnxExporterError
|
If the |
Examples:
>>> import tempfile
>>> import torch
>>> from pathlib import Path
>>> from lucid_yolo.models.build import Detector
>>> deployed = Detector("n", num_classes=4).eval().deploy().eval()
>>> graph = DetectExportGraph(deployed, canvas=(64, 64)).eval()
>>> with tempfile.TemporaryDirectory() as tmp:
... written = export_graph(graph, Path(tmp) / "detect.onnx")
... written.is_file()
True
Graphs¶
Bases: Module
A deploy() view plus its E2E decode, as one exportable module.
The shared base of :class:DetectExportGraph, :class:SegmentExportGraph, :class:OrientedExportGraph and
:class:KeypointExportGraph: it bakes the anchor grid for one canvas into buffers and holds the
:class:~lucid_yolo.decode.topk_e2e.TopKDecoder whose k fixes every task's output length — the axis-aligned
tasks decode with it directly, the oriented task with a
:class:~lucid_yolo.decode.topk_e2e.RotatedTopKDecoder built at the same k — so a subclass need only
compose its own task's dense outputs into a decoded tuple. forward is left to subclasses — that is where
the four tasks differ, per the module docstring's compositions — so this class alone is not directly exportable.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
deployed
|
Module
|
The |
required |
canvas
|
tuple[int, int]
|
Input canvas |
required |
device
|
device | None
|
Device the anchor buffers are built on and |
None
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If either |
Examples:
>>> import torch
>>> from lucid_yolo.models.build import Detector
>>> deployed = Detector("n", num_classes=4).eval().deploy().eval()
>>> graph = E2EExportGraph(deployed, canvas=(128, 128))
>>> graph.anchor_points.shape, graph.strides.shape # 16^2 + 8^2 + 4^2 = 336
(torch.Size([336, 2]), torch.Size([336]))
Bases: E2EExportGraph
Detection E2E export graph: dense one-to-one outputs to the A9 tuple.
Reproduces :func:~lucid_yolo.predict.predict_image's "e2e" composition — the deployed model's one-to-one
(cls, box) pair straight into :class:~lucid_yolo.decode.topk_e2e.TopKDecoder, same argument order — minus
the letterbox inverse and the confidence filter, which are single-image concerns a batched, fixed-shape graph
does not have.
Examples:
>>> import torch
>>> from lucid_yolo.models.build import Detector
>>> deployed = Detector("n", num_classes=4).eval().deploy().eval()
>>> graph = DetectExportGraph(deployed, canvas=(128, 128)).eval()
>>> with torch.no_grad():
... detections = graph(torch.zeros(1, 3, 128, 128))
>>> detections.shape # (B, k, 6): the A9 tuple, k = 300 (A9)
torch.Size([1, 300, 6])
Decode an image batch into fixed-size A9 detections.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Detections of shape |
Tensor
|
|
Tensor
|
meth: |
Bases: E2EExportGraph
Segmentation E2E export graph: the A9 tuple beside its Eq. 7 instance masks.
Reproduces :func:~lucid_yolo.predict.predict_segmentation's "e2e" composition — the same decoder call for
the boxes, the same :meth:~lucid_yolo.decode.topk_e2e.TopKDecoder.decode_with_indices source of the anchor
indices mask coefficients are gathered by, and the same mask assembly, resize, crop and threshold operations.
Export retains the original single-resize graph while eager decode chunks its full-canvas temporaries. One shape
divergence remains: every one of the k rows is gathered, padding rows included, because the output shape is
fixed.
That divergence is why this graph's supported canvases are narrower than its base class's: a canvas yielding
fewer anchors than k is refused in __init__, where it can still be reported as a canvas choice. The
coefficient gather is the shared clamp-and-zero one, so such a canvas would no longer fail at trace time; the
module docstring states why the refusal is kept anyway.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
deployed
|
Module
|
The |
required |
canvas
|
tuple[int, int]
|
Input canvas |
required |
device
|
device | None
|
Device the buffers are built on and |
None
|
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
Examples:
>>> import torch
>>> from lucid_yolo.models.build import Segmenter
>>> deployed = Segmenter("n", num_classes=4).eval().deploy().eval()
>>> graph = SegmentExportGraph(deployed, canvas=(128, 128)).eval()
>>> with torch.no_grad():
... detections, masks = graph(torch.zeros(1, 3, 128, 128))
>>> detections.shape, masks.shape # (B, k, 6) beside (B, k, H, W)
(torch.Size([1, 300, 6]), torch.Size([1, 300, 128, 128]))
>>> SegmentExportGraph(deployed, canvas=(64, 64)) # 84 anchors, below k
Traceback (most recent call last):
...
ValueError: canvas (64, 64) yields 84 anchors, below the 300-detection cap; padding rows have no coefficients
Decode an image batch into fixed-size A9 detections and their masks.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
A pair of the |
Tensor
|
|
tuple[Tensor, Tensor]
|
(see :func: |
tuple[Tensor, Tensor]
|
row-aligned. |
Bases: E2EExportGraph
Oriented E2E export graph: dense one-to-one outputs to the A45 tuple.
Reproduces :func:~lucid_yolo.predict.predict_oriented's "e2e" composition — one
:class:~lucid_yolo.decode.topk_e2e.RotatedTopKDecoder call over the deployed model's one-to-one class logits,
(box, angle) pair and the baked anchor grid, at the same k as the other graphs' decoder.
Examples:
>>> import torch
>>> from lucid_yolo.models.build import OrientedDetector
>>> deployed = OrientedDetector("n", num_classes=4).eval().deploy().eval()
>>> graph = OrientedExportGraph(deployed, canvas=(128, 128)).eval()
>>> with torch.no_grad():
... detections = graph(torch.zeros(1, 3, 128, 128))
>>> detections.shape # (B, k, 7): the A45 tuple, k = 300 (A9)
torch.Size([1, 300, 7])
Decode an image batch into fixed-size A45 rotated detections.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
Detections of shape |
Tensor
|
|
Tensor
|
meth: |
Bases: E2EExportGraph
Keypoint E2E export graph: the A9 tuple beside each detection's point set.
The fourth task's graph, and the last of the four to exist -- 0.5.0 shipped a keypoint model that trains, validates and evaluates but had no exported view at all, so it was the only released task whose model could not leave PyTorch (WP-151).
The composition follows :class:SegmentExportGraph rather than :class:DetectExportGraph, for the reason
segmentation has it: the point stem is dense over anchors, so the decode needs the source anchor of each
detection row, which is what :meth:~lucid_yolo.decode.topk_e2e.TopKDecoder.decode_with_indices returns and
plain :meth:~lucid_yolo.decode.topk_e2e.TopKDecoder.forward discards. The ordering is the one
WP-062's angle branch established and :class:~lucid_yolo.ptl.module.DetectionLitModule already uses for its own
val/oks_mAP: decode the whole dense tensor first, gather second.
:func:~lucid_yolo.models.heads.keypoint.decode_keypoints needs the entire anchor grid and offers no
per-detection indexing, so gathering first would mean reconstructing each kept row's anchor and stride by hand.
Padding rows are handled by :func:~lucid_yolo.decode.common.gather_by_anchor, which clamps a padding index to
a valid position and then zeroes that row, so a padding row yields the origin rather than the last anchor's
pose. A fixed-shape graph cannot drop rows, and this task's gather is defined on them, so this graph accepts
every canvas its base class does; :class:SegmentExportGraph uses the same gather but refuses canvases that
would produce padding rows.
Examples:
>>> import torch
>>> from lucid_yolo.models.build import KeypointDetector
>>> deployed = KeypointDetector("n", num_classes=4, num_keypoints=3).eval().deploy().eval()
>>> graph = KeypointExportGraph(deployed, canvas=(128, 128)).eval()
>>> with torch.no_grad():
... detections, keypoints = graph(torch.zeros(1, 3, 128, 128))
>>> detections.shape, keypoints.shape # (B, k, 6) beside (B, k, K, 2)
(torch.Size([1, 300, 6]), torch.Size([1, 300, 3, 2]))
Decode an image batch into fixed-size A9 detections and their point sets.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
A pair of the |
Tensor
|
point coordinates in letterboxed-canvas pixels, row-aligned. Padding |
tuple[Tensor, Tensor]
|
rows carry the origin and score zero. |