🏋️ Training

The Lightning pieces a training run is assembled from: the module that owns the model, loss and optimizer, the data module that builds the loaders, and the two callbacks the recipe needs. Generated from the docstrings in src/lucid_yolo/ptl/.

Module and data

DetectionLitModule

Bases: LightningModule

Lightning detector: model stack, dual-branch loss, and MuSGD (WP-034).

Composes :class:~lucid_yolo.models.backbone.DetectionBackbone, :class:~lucid_yolo.models.neck.DetectionNeck, and :class:~lucid_yolo.models.heads.detect.DualDetectionHead through :func:~lucid_yolo.models.build.build_detection_stages — the same factory :class:~lucid_yolo.models.build.Detector uses, so the model the fidelity gate measures and the model training runs cannot drift apart (WP-087). It scores their dense predictions with :class:~lucid_yolo.losses.dual_loss.DualBranchLoss, and optimizes under Lightning automatic optimization (D4) with :class:~lucid_yolo.optim.musgd.MuSGD. See the module docstring for the A8 constant-LR deferral, the task conditioning, and the :attr:alpha seam.

Parameters:

Name	Type	Description	Default
depth	float	Depth multiplier scaling per-stage repeat counts of the backbone and neck.	required
width	float	Width multiplier scaling channel counts.	required
max_channels	int	Channel cap applied before the width multiply.	required
num_classes	int	Number of object classes the head predicts.	required
task	str	Supervision task; one of "detect", "segment", "obb" or "keypoints", all four active. "segment" additionally builds the head's mask-coefficient stems and the three :func:~lucid_yolo.models.build.build_segmentation_stages branches as proto_fusion/protonet/semantic, and supervises them with the two mask terms below. "obb" instead builds both head branches' orientation stems (predict_angle, A20) and swaps the two box terms for their rotated counterparts (see the module docstring's table). "keypoints" builds both branches' point stems (num_keypoints, WP-122) and adds the R14 residual log-likelihood term at keypoint_gain and the OKS term at oks_gain, leaving every detection term as it was. Defaults to "detect".	'detect'
num_keypoints	int | None	Point count K both head branches predict under task="keypoints"; required for that task and refused for every other one. It has no default on purpose — K describes the dataset's annotation schema, not the method (the module docstring argues the case). The two directions are both errors on purpose: a missing count under "keypoints" has nothing to guess, and a count under any other task names a schema the head builds no stem for, so the module would train no points while save_hyperparameters recorded K for a reader to trust later.	None
lr	float	Base learning rate for MuSGD (lr0; the A8 schedule decays from it). Defaults to 0.01.	0.01
lrf	float	Final LR fraction of the A8 linear decay — the LR ends at lr * lrf. >= 1 together with warmup_epochs <= 0 disables the schedule (constant LR). Defaults to 0.01.	0.01
warmup_epochs	float	Length of the opening linear LR warmup, in epochs (fractions allowed). 0 disables warmup. Defaults to 3.0.	3.0
momentum	float	MuSGD momentum coefficient. Defaults to 0.95.	0.95
weight_decay	float	Decoupled weight decay (matrix parameters only). Defaults to 5e-4.	0.0005
w_muon	float	Additive gain on the MuSGD Muon branch. Defaults to 0.5.	0.5
w_sgd	float	Additive gain on the MuSGD SGD branch. Defaults to 0.5.	0.5
box_gain	float	Gain on the IoU term of both loss branches — the axis-aligned Complete-IoU under "detect"/"segment", and under "obb" the rotated ProbIoU that replaces it in the same slot (A49; the module docstring explains why the slot rather than the formula is what the name refers to). Defaults to 7.5.	7.5
cls_gain	float	Classification-term gain shared by both branches. Defaults to 0.5.	0.5
l1_gain	float	Gain on the L1 box term of both branches: against the axis-aligned target under "detect"/"segment", and under "obb" against the rotated box's own (cx, cy, w, h) (A50). Defaults to 6.0.	6.0
alpha	float	Initial one-to-many branch weight seeding the dual loss (used before training and whenever the epoch schedule is inactive). Defaults to 0.5.	0.5
alpha_init	float	One-to-many branch weight the progressive schedule sets on the first epoch (branch weights (0.8, 0.2)). Defaults to 0.8.	0.8
alpha_final	float	One-to-many branch weight the schedule ramps to on the last epoch (branch weights (0.1, 0.9)). Defaults to 0.1.	0.1
mask_gain	float	Weight on the instance-mask term under task="segment", ignored otherwise (A38). Defaults to 2.5.	2.5
semantic_gain	float	Weight on the auxiliary semantic term under task="segment", ignored otherwise (A38). Defaults to 0.5.	0.5
angle_gain	float	Weight on R1 Eq. 15's square-object angle term under task="obb", ignored otherwise. Defaults to 0.25 (A22 — R1 does not state one, and :func:~lucid_yolo.losses.angle_loss.square_angle_loss returns the term pre-gain precisely so this caller owns it). Lowered from 1.0 by WP-093, which measured the term destabilizing angle regression on elongated targets: at 1.0 the oriented overfit cleared its floor on 1 of 5 seeds with a 0.667 spread, at 0.25 on 4 of 5 with 0.096. See A22 for the dose-response and why the remedy is this gain rather than lambda, which R1 does state.	0.25
rotated_iou_form	str	Which of R17's two rotated-IoU losses the "obb" box term uses, "hellinger" (bounded, the A49 default) or "bhattacharyya" (unbounded). Ignored otherwise.	DEFAULT_ROTATED_IOU_FORM
keypoint_gain	float	Weight on the R14 residual-log-likelihood term under task="keypoints", ignored otherwise. Defaults to 1.0 (A68), which is not a measured value: R14 trains RLE as an objective in its own right and never states a weight for it beside a detection loss, and no allowlisted source covers the pair. 1.0 is therefore the neutral placeholder — the term as the paper writes it — pending the same dose-response treatment WP-093 gave angle_gain under A22. Expect it to want lowering, and by orders rather than factors: the term is a negative log-likelihood of residuals standardized by a sigmoid sigma in (0, 1), while :func:~lucid_yolo.models.heads.keypoint.decode_keypoints works in absolute input pixels, so at initialisation it is measured two orders above the whole detection total rather than beside it. A68 cites this module's keypoint_gain=0 bit-exactness test as its validation: at zero the term must leave the detection objective untouched, whatever scale it takes at one.	1.0
keypoint_loss	str	Which keypoint objective task="keypoints" supervises with, ignored otherwise. Defaults to "rle", R14's full residual log-likelihood (:class:~lucid_yolo.losses.rle_loss.RLELoss) — every config, checkpoint and accepted figure this project has produced was produced under it, and the default is what keeps that true. The alternative, "laplace_nll" (:class:~lucid_yolo.losses.keypoint_nll_loss.LaplaceNLLLoss), drops the normalizing flow and keeps everything else, which is R14 Table 7's own published ablation: the two rows are "RLE" at 70.5 AP and "Laplace, learnable variance" at 67.4 AP on COCO. It exists for WP-125, whose acceptance is RLE's mechanism claim rather than an absolute pose number — a claim of that shape needs a paired run differing only in the mechanism, and one figure on its own cannot settle it. The switch is therefore an experimental control, not a tuning knob: a production pose run has no reason to leave the default.	'rle'
oks_gain	float	Weight on the Object-Keypoint-Similarity term (:class:~lucid_yolo.losses.oks_loss.OKSLoss, WP-178) under task="keypoints", ignored otherwise. Defaults to 0.0 (A75), which is a deliberate off rather than a measured value. R1 sec. 3.4.2 states that YOLO26's pose training uses an OKS-based loss extended with RLE, and R1 Table 9 ablates the pair's weights (w_OKS, w_RLE) on COCO keypoints val for YOLO26s — (48, 0) 61.5, (48, 1) 60.8, (24, 1) 63.0, (24, 2) 62.5, (12, 1) 62.6, (0, 1) 61.9, (0, 2) 62.4 mAP — so the published target is (24, 1). But R1 gives no formula for its OKS loss and its citation is off this project's allowlist, so the term here is derived from R12's metric definition (A75), and a weight measured under a loss this project had to derive is not a measured value here. Shipping 24.0 as the default would make an untested number look like evidence while being none — A68's own reasoning for keypoint_gain — and would silently move every existing keypoints config off the objective its accepted figures were produced under. 0.0 keeps that objective bit-exact (the module's oks_gain=0 bit-exactness test is its validation); the pre-gain term is still computed and logged as train/oks so a dose-response run against (24, 1) has its reading before it moves the gain.	0.0
fused_assignment	bool	Opt into CUDA target-assignment fusion. It retains the current objectives and native CPU/MPS, compiler and unsupported-layout paths. Default False avoids compilation cost on short or small workloads. Where compile_step=True is honoured, the compiled objective assigns natively and :meth:setup warns once per fit.	False
compile_step	bool	Whether to :func:torch.compile the detection objective — forward, both decodes, assignment and the dual loss as one Inductor region (WP-184, D24). Defaults to False. Honoured on a CUDA accelerator under task="detect" only; any other accelerator or task warns once at :meth:setup and stays eager (MPS Inductor crashes in convolution_backward, and the three other tasks' extra terms still read tensors on the host). Measured on an RTX PRO 6000 at variant n, COCO, batch 16, deterministic: the training step goes from 74.7 ms to 60.9 ms (x1.23) against eager, for 110-140 s of compilation per process plus a handful of recompiles per run (the no_grad evaluation graph, and one each for a batch with no ground truths, one per image, or one image — Dynamo specialises sizes 0 and 1 whatever is marked). It compiles the bound method rather than the module, so parameter names, the state dict, EMA, checkpoints and export are untouched. Whether the compiled objective trains to the same mAP as eager is A78, still open — which is why the default is off.	False
capture_orthogonalization	bool	Enable MuSGD's optional CUDA graph replay of its existing Newton--Schulz calls. Defaults to False. Shapes, strides, dtype, device, step count and precision select each recording; this flag alone leaves learning rates, momentum and parameters outside the graphs. Extra workspace and setup can outweigh the saving on short runs.	False
capture_network	bool	Replay fixed detection, keypoint, oriented or segmentation network stages' CUDA forward and backward with :func:torch.cuda.make_graphed_callables. Defaults to False because capture pays setup and retains a graph pool. Requires task="detect", task="keypoints", task="obb" or task="segment", compile_step=False and single-device CUDA automatic optimization. Pose records backbone, neck and each prediction stem independently, preserving absent point-stem gradients when a batch supplies no positives. Pose training stays eager with a fully frozen backbone and neck, or when the head's topology or class forward changed since construction or a head module is hooked, overridden or compiled. The trainable set is fixed per fit at the first recording: a fully frozen start, or a later freeze or unfreeze of a recorded parameter, keeps the rest of that fit eager after one warning. Each pose stem separately retires its recording if its trainable set or mode changes, and replacing a head module retires every recording before shared features run. Oriented training records the complete network, including both angle stems, and retains the native classification and oriented objective. An overridden objective keeps it eager. Detection, pose and oriented training stay eager when the LightningModule carries a hook, overrides forward or is compiled. Segmentation records only backbone and neck, leaving all predictions independent; capture_segment_heads can replay those predictions separately. Segmentation retires replay when the head topology/dispatch changes, a head hook is registered, or forward_segmentation is overridden or compiled. Assignment, losses, schedules, optimizer, EMA and evaluation stay eager. One image signature is recorded per fit; partial batches stay eager and changed storage, modes, hooks or precision retire the recording. Warmup restores buffers and RNG, and the fit callbacks discard recordings on completion or failure. Higher-order differentiation is unsupported.	False
capture_updates	bool	Enable MuSGD's optional replay of the existing update preparation, including momentum, followed by a native batched parameter write at the current learning rate. Defaults to False. Uses private gradient storage and live parameters and momentum; unsupported signatures stay native. Like capture_orthogonalization, this has setup and retained-workspace costs.	False
capture_loss	bool	Replay the existing detection assignment and loss forward/backward, for detection or pose, on single-device CUDA automatic training. Default off: recording pays setup and retains a graph pool. Masked target capacity grows geometrically; the method, label checks, live branch schedule and optimizer stay unchanged. Requires task="detect" or task="keypoints", and compile_step=False. Pose also batches the existing RLE densities, keeping visibility, separate means and branch weights. Batched parameter-gradient summation can change rounding. Can combine with capture_network, capture_orthogonalization and capture_updates; it records no optimizer update itself, so the optimizer step is whatever those optimizer options select.	False
capture_segment_heads	bool	Record independent CUDA segmentation prediction groups, including first-order backward. Defaults to False. Requires task="segment" and compile_step=False, PyTorch 2.9 or newer, BF16 mixed precision and single-device automatic optimization; gradient accumulation replays too. CPU, MPS, FP32, FP16 and validation keep native execution. compile_step=True is refused rather than ignored, although it leaves the segment objective eager today, so an accepted configuration keeps its meaning once that is compiled.	False
capture_gradient_norm	bool	Replay PyTorch's L2 gradient-norm reduction on single-device CUDA with torch>=2.9. Default false. Native scaling reads the current clipping limit; one fit-scoped gradient signature is retained and drift retires recording. Value clipping, disabled clipping (limit None or at most zero), optimizers that unscale inside their own step, customized module/precision-plugin or PyTorch clipping, non-CUDA devices, unsupported torch versions and unsupported gradient layouts keep native clipping. Ordinary 16-mixed GradScaler training can capture after unscaling; unscale-in-step optimizers retain Lightning's mixed-precision clipping rejection. Extra private CUDA buffers and setup cost may outweigh the gain outside small-batch training.	False

Raises:

Type

Description

ValueError

If task is not one of "detect", "segment", "obb", "keypoints", if rotated_iou_form is not a known rotated-IoU form, if capture_network is combined with a task other than detection, keypoints, oriented detection or segmentation, capture_loss with a task other than detection or keypoints, capture_segment_heads with a non-segmentation task, or any of them with compile_step=True (for capture_segment_heads, forward compatibility with a compiled segment objective), if keypoint_loss is not a known keypoint objective, if task="keypoints" was asked for without a num_keypoints, or if a num_keypoints was given for a task that builds no point stem.

Examples:

>>> import torch
>>> from lucid_yolo.data.targets import Targets
>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4)
>>> images = torch.zeros(1, 3, 160, 160)
>>> targets = [Targets(boxes=torch.tensor([[8.0, 8.0, 40.0, 40.0]]), labels=torch.tensor([1]))]
>>> loss = module.training_step((images, targets), 0)  # doctest: +SKIP
>>> bool(torch.isfinite(loss))  # doctest: +SKIP
True

task property

task: str

Supervision task this module was built for: "detect", "segment", "obb" or "keypoints".

Read-only, and the supported way for a consumer to ask whether a loaded checkpoint has a mask branch, an angle branch or a keypoint branch — :meth:forward_segmentation gates on this same value, so the caller's question and the module's own behaviour cannot answer differently. Reading hparams["task"] instead would be a stringly-typed lookup into a bag that is only as complete as the checkpoint that filled it, and it raises rather than defaulting when a checkpoint predates the hyperparameter.

Returns:

Type	Description
str	The task string given at construction.

Examples:

>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4)
>>> module.task
'detect'

alpha property writable

alpha: float

One-to-many branch weight of the dual loss (the WP-035 schedule seam).

Returns:

Type	Description
float	The current :attr:DualBranchLoss.alpha; setting it rewrites the underlying loss attribute so a scheduler
float	can update the ramp per epoch.

Examples:

>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4)
>>> module.alpha = 0.8
>>> module.loss.alpha
0.8

save_hyperparameters

save_hyperparameters(*args: Any, ignore: Sequence[str] | str | None = None, frame: FrameType | None = None, logger: bool = True) -> None

Record Lightning hyperparameters, with task in its plain spelling, on every save.

The plain spelling, even when a :class:~lucid_yolo.tasks.Task member was passed: a checkpoint loads under torch.load(weights_only=True), whose restricted unpickler refuses a Task in hparams["task"]. Normalizing the constructor's local before the call is not enough, because the save merges every subclass __init__ frame's arguments over this one's, **kwargs contents included, and correcting the record once after the constructor's own call is not enough either, because a subclass may save again after super().__init__ returns. So every save corrects both records it wrote: hparams, which the checkpoint stores, and the initial snapshot a logger writes. The arguments keep Lightning's meaning; with no frame, the caller's frame is passed on, since Lightning would otherwise inspect this override's. A hyperparameter dict given without task (LightningCLI's path) stays as given.

Parameters:

Name	Type	Description	Default
*args	Any	Lightning's positional hyperparameters: argument names, or one dict or namespace.	()
ignore	Sequence[str] | str | None	Argument names Lightning leaves out of the record.	None
frame	FrameType | None	The frame whose __init__ arguments are recorded; the caller's when None.	None
logger	bool	Whether the initial snapshot is offered to the logger.	True

Examples:

>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=64, num_classes=3)
>>> module.save_hyperparameters({"task": Task.DETECT}, logger=False)
>>> type(module.hparams["task"]) is str
True

setup

setup(stage: str) -> None

Pick the CUDA memory layout and objective, then refuse multi-process runs the protocols cannot gather.

With compile_step on, :meth:_compile_objective decides once whether the detection objective is compiled (CUDA accelerator, "detect" task) or stays eager with a warning naming why; a second setup — fit then validate — changes nothing.

On a CUDA accelerator the parameters are converted to channels_last (NHWC), the layout cuDNN's convolution kernels prefer: eager training of the n model stepped 8% faster on an RTX PRO 6000 with no other change, and the datamodule hands over its images in the same layout so no op has to permute. This hook is where the conversion belongs because it runs before device placement and before every callback's on_fit_start — so :class:EMACallback's shadow, cloned there, inherits the layout rather than averaging across a permute. The decision reads the trainer's accelerator rather than :func:torch.cuda.is_available, which is true on a CUDA box that was asked for accelerator="cpu", and it is not a flag: CPU and MPS keep their default layout and their bit-exact goldens, and there is nothing to configure on the one backend where the layout pays. A module with no trainer attached is left as built.

All task metrics gather whole-epoch records across ranks before scoring. Rotated and keypoint metrics also drop sampler padding by dataset identity, retaining the original sequential validation order and scoring protocols.

Setting up "fit" also re-arms the once-per-fit gradient-snapshot-skip verdict that :meth:_report_native_closure logs at the first training step.

Parameters:

Name	Type	Description	Default
stage	str	The Trainer stage being set up.	required

Examples:

>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4)
>>> module.setup("fit")  # "detect" scales freely; no trainer attachment needed

configure_callbacks

configure_callbacks() -> list[Callback]

Keep opted-in CUDA recordings scoped to one fit.

configure_gradient_clipping

configure_gradient_clipping(optimizer: Optimizer, gradient_clip_val: int | float | None = None, gradient_clip_algorithm: str | None = None) -> None

Optionally record the norm reduction while retaining native clipping semantics.

Ordinary FP16 GradScaler gradients are already unscaled before this hook and can use capture. Optimizers that unscale inside their own step keep the native call, including Lightning's mixed-precision clipping rejection. Value or disabled clipping retires capture once, as does customized module or precision-plugin clipping, with the latter diagnostic naming the plugin type. Plugins inheriting the checked stock clipping methods through their MRO remain eligible.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Optimizer whose gradients are clipped.	required
gradient_clip_val	int | float | None	Clipping limit; None uses the trainer's setting. Limits at or below zero disable clipping.	None
gradient_clip_algorithm	str | None	Clipping algorithm; None uses the trainer's setting. Algorithms other than "norm" keep native clipping.	None

on_train_epoch_start

on_train_epoch_start() -> None

Ramp the dual-loss branch weight for the epoch about to start (WP-035).

Sets :attr:alpha to :meth:~lucid_yolo.losses.progressive.ProgressiveLossSchedule.alpha_at evaluated at the current 0-based epoch and the trainer's total epoch count — the linear R1 Eq. 2-3 ramp updated once per epoch. When trainer.max_epochs is unset (None or <= 0, as for a step-bounded or fast_dev_run run) the ramp denominator is undefined, so alpha is left at its seeded value.

Examples:

>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4)
>>> module.alpha  # seeded value before any epoch starts
0.5

forward

forward(images: Tensor) -> DualHeadOutput

Run the backbone, neck, and dual head over an image batch.

Parameters:

Name	Type	Description	Default
images	Tensor	Input batch of shape (B, 3, H, W) with H and W divisible by 32.	required

Returns:

Name	Type	Description
The	DualHeadOutput	class:~lucid_yolo.models.heads.detect.DualHeadOutput dense predictions (raw class logits and raw
	DualHeadOutput	ltrb distances) of both branches.

Examples:

>>> import torch
>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4).eval()
>>> with torch.no_grad():
...     out = module(torch.zeros(1, 3, 160, 160))
>>> out.o2o_cls.shape[-1]
4

forward_segmentation

forward_segmentation(images: Tensor) -> SegmentOutput

Run the detection and mask branches over an image batch (WP-087).

:meth:forward returns the detection head's output alone, because that is what every detection consumer — the E2E decode, the mAP metric, the export path — asks for. Segmentation needs the prototype maps and the auxiliary logits from the same neck features, so this is the second entry point, and it is the module's only composition of the mask side: the training step and the overfit gate both call it rather than re-running backbone -> neck -> proto_fusion -> protonet themselves, which is how the two would drift apart.

The auxiliary semantic branch is training-only (A17) and returns None in eval mode; the prototypes are produced in both modes.

Parameters:

Name	Type	Description	Default
images	Tensor	Input batch of shape (B, 3, H, W) with H and W divisible by 32.	required

Returns:

Name	Type	Description
A	SegmentOutput	class:~lucid_yolo.models.build.SegmentOutput holding the dual head's dense predictions (including both
	SegmentOutput	branches' mask coefficients), the raw prototype maps, and the auxiliary logits.

Raises:

Type	Description
ValueError	If this module's task is not "segment" — the mask branches are only built for that task, so any other task has no prototypes to return.

Examples:

>>> import torch
>>> module = DetectionLitModule(
...     depth=0.34, width=0.25, max_channels=256, num_classes=4, task="segment"
... ).eval()
>>> with torch.no_grad():
...     out = module.forward_segmentation(torch.zeros(1, 3, 64, 64))
>>> out.prototypes.shape  # twice the P3 grid (A15): 64 / 8 * 2
torch.Size([1, 32, 16, 16])
>>> out.detect.o2o_coeff.shape[-1], out.semantic is None  # aux is training-only (A17)
(32, True)

forward_branch

forward_branch(images: Tensor, branch: BranchName) -> BranchOutput

Run the backbone, neck, and one head branch over an image batch.

:meth:forward returns both branches because training reads both. Every inference decode reads exactly one: the top-k path the one-to-one branch, the suppression path the one-to-many. Called through :meth:forward, the other branch's stems run over all three levels and their outputs are discarded, so this is the entry point for a caller that already knows which half it will read.

The composition is deliberately the same three calls :meth:forward makes, with only the head's own branch selection narrowed -- the backbone and neck are shared by both branches and are not run twice or run differently here.

Parameters:

Name	Type	Description	Default
images	Tensor	Input batch of shape (B, 3, H, W) with H and W divisible by 32.	required
branch	BranchName	Which head branch to run, one of :data:~lucid_yolo.models.heads.detect.BRANCH_NAMES.	required

Returns:

Type	Description
BranchOutput	That branch's :class:~lucid_yolo.models.heads.detect.BranchOutput dense predictions.

Raises:

Type	Description
ValueError	If branch is not a known branch name.

Examples:

>>> import torch
>>> module = DetectionLitModule(depth=0.34, width=0.25, max_channels=1024, num_classes=4).eval()
>>> with torch.no_grad():
...     one_to_one = module.forward_branch(torch.zeros(1, 3, 160, 160), "o2o")
>>> one_to_one.cls.shape[-1]
4

forward_segmentation_branch

forward_segmentation_branch(images: Tensor, branch: BranchName) -> tuple[BranchOutput, Tensor]

Run one head branch and the prototype maps over an image batch.

The branch-limited twin of :meth:forward_segmentation, and the mask side's reason for existing is the same as the detection side's: Eq. 7 reads the coefficients of the branch its decoder read, never both.

Two things are dropped rather than returned. The other branch's stems -- including its coefficient stem, which is the widest of the optional ones -- are never run. The auxiliary semantic head is not run either: it is training-only (A17) and already returns None in eval mode, so a caller decoding masks has nothing to read from it. What remains is exactly the pair Eq. 7 multiplies.

Parameters:

Name	Type	Description	Default
images	Tensor	Input batch of shape (B, 3, H, W) with H and W divisible by 32.	required
branch	BranchName	Which head branch to run, one of :data:~lucid_yolo.models.heads.detect.BRANCH_NAMES.	required

Returns:

Type	Description
BranchOutput	A (branch_output, prototypes) pair: that branch's dense predictions, and
Tensor	the raw prototype maps (B, num_coeffs, H/4, W/4).

Raises:

Type	Description
ValueError	If this module's task is not "segment" -- the mask branches are only built for that task, so any other task has no prototypes to return; or if branch is not a known branch name.

Examples:

>>> import torch
>>> module = DetectionLitModule(
...     depth=0.34, width=0.25, max_channels=256, num_classes=4, task="segment"
... ).eval()
>>> with torch.no_grad():
...     one_to_one, prototypes = module.forward_segmentation_branch(torch.zeros(1, 3, 64, 64), "o2o")
>>> prototypes.shape  # twice the P3 grid (A15): 64 / 8 * 2
torch.Size([1, 32, 16, 16])
>>> one_to_one.coeff.shape[-1]
32

training_step

training_step(batch: StepBatch, batch_idx: int) -> Tensor

Run one training step under automatic optimization (D4).

Parameters:

Name	Type	Description	Default
batch	StepBatch	The datamodule batch (images, list[Targets], masks); the two-element (images, list[Targets]) form rasterises in the step.	required
batch_idx	int	Index of the batch within the epoch (unused).	required

Returns:

Type	Description
Tensor	The scalar total loss for Lightning to backpropagate.

optimizer_step

optimizer_step(epoch: int, batch_idx: int, optimizer: Optimizer | LightningOptimizer, optimizer_closure: Callable[[], Any] | None = None) -> None

Omit snapshots that the native closure clears before backward.

Only detection network replay with accumulation one and the native pass-through profiler qualifies. The native forward runs inside the permission; zeroing, backward and subsequent callbacks run outside it. Restoring the closure in finally keeps a failed or reused closure from retaining the wrapper. Custom zeroing, observers, strategies, precision plugins, optimizer wrappers and a closure whose own closure method was replaced on the instance keep the original snapshots. Lightning minors beyond the inspected 2.4-2.6 contract also retain them. This hook only filters on the replay flag and the closure's shape; the certificate, evaluated as the forward starts, decides. The first step of a fit logs the verdict once, at rank-zero info level: active, or the first check that failed.

validation_step

validation_step(batch: StepBatch, batch_idx: int) -> Tensor

Run one validation step: shared forward and loss, plus the mAP update.

Beyond the val/-logged loss, the one-to-one branch is decoded with the E2E :class:~lucid_yolo.decode.topk_e2e.TopKDecoder from the same forward and accumulated into the epoch's :class:~torchmetrics.detection.MeanAveragePrecision (WP-077), logged as val/mAP by :meth:on_validation_epoch_end. Scoring runs in letterbox coordinates — IoU is invariant to each image's uniform letterbox scaling, so the number tracks the original-coordinate protocol closely; the acceptance figure remains lucid-eval (original coordinates, both paths).

A "segment" module additionally decodes the kept detections' masks and accumulates val/segm_mAP (:meth:_update_val_segm), so a segmentation run reports the quality of the branch it exists for rather than of its boxes alone. That needs the ground-truth masks, which only a loader built with mask_targets=True supplies — the CLI links that flag to model.task, and a batch without them is scored on boxes only rather than rasterised a second time here.

An "obb" module accumulates the WP-063 rotated mAP (:meth:_update_val_rotated) instead of val/mAP, not beside it (WP-102): the axis-aligned figure reads the A44 composition's pre-rotation rectangle, so a run whose orientations were random and one whose orientations were right log the identical curve, and the acceptance number would otherwise first appear hours after the run.

A "keypoints" module additionally decodes the kept detections' points and accumulates val/oks_mAP (:meth:_update_val_keypoints, WP-137), the same "beside val/mAP" relationship "segment" has: the box branch and the point branch answer different questions, and a run whose boxes are fine but whose points have not yet wired up should not read as healthy on the one number that was already there.

Parameters:

Name	Type	Description	Default
batch	StepBatch	The datamodule batch (images, list[Targets], masks); the two-element (images, list[Targets]) form rasterises in the step.	required
batch_idx	int	Index of the batch within the epoch (unused).	required

Returns:

Type	Description
Tensor	The scalar total validation loss.

on_validation_epoch_end

on_validation_epoch_end() -> None

Compute and log the epoch's E2E val/mAP (progress-bar metric), then reset.

A "segment" run additionally logs val/segm_mAP, whenever any batch of the epoch carried ground-truth masks; an "obb" run logs the WP-063 val/rotated_mAP and val/rotated_mAP50 over everything the epoch saw and logs no val/mAP at all (WP-102): that figure reads the A44 composition's pre-rotation rectangle, so it cannot separate a run whose orientations were right from one whose orientations were random, and its accumulator costs a CPU pass over every detection of every batch to say so. The rotated metric is the one an oriented run is asking for; carrying a second one that answers a different question is not a free comparison.

A "keypoints" run additionally logs val/oks_mAP (WP-137, :meth:_log_oks_map), beside val/mAP rather than instead of it, for the same reason val/segm_mAP sits beside val/mAP and not the other way val/rotated_mAP does.

The val/segm_mAP guard is reduced across processes before it is read. :meth:~torchmetrics.Metric.compute synchronises collectively, so a rank that happened to draw a shard without mask targets must still enter it — deciding to skip on that rank alone would leave the others waiting inside the collective for a call that never comes. The reduction is any: one rank with masks makes the metric well-defined for all of them, and on a single process it is the identity, so the accepted single-device behaviour is unchanged.

configure_optimizers

configure_optimizers() -> OptimizerLRScheduler

Return MuSGD, paired with the A8 warmup + linear-decay LR schedule.

The schedule (WP-072) is a per-step :class:~torch.optim.lr_scheduler.LambdaLR over :func:~lucid_yolo.optim.schedule.warmup_decay_factor: a linear warmup across the first warmup_epochs epochs, then a linear decay from lr down to lr * lrf at the trainer's estimated final step. The bare constant-LR optimizer is returned instead when the schedule is explicitly disabled (lrf >= 1 and warmup_epochs <= 0), when no trainer is attached (direct calls in tests and tools), or when the run has no positive max_epochs to anchor the warmup fraction.

The warmup length is a fraction of the run, so a short run can ask for more warmup steps than it has steps: warmup_epochs=3 against a two-epoch smoke budget is config-reachable and is exactly the shape the short tiers get launched at. Left alone, every step would take the ramp branch — the LR would climb for the whole run, peak below lr, and never decay, with nothing saying so. The count is clamped to total_steps - 1, which recovers the peak but not the decay: a run whose requested warmup covers it has no steps left to decay over, so the clamp buys back only the ramp reaching full lr on the final step. That is the ceiling of what a clamp can do here, which is why it warns rather than applying silently — the recipe cannot be honoured, and the honest signal is to say which schedule actually ran. Clamping rather than raising keeps the launchable run launchable: a crash here would fail a run over a schedule detail on the tier where the schedule matters least.

Returns:

Name	Type	Description
A	OptimizerLRScheduler	class:~lucid_yolo.optim.musgd.MuSGD over self.parameters(), alone or inside a Lightning
	OptimizerLRScheduler	optimizer/scheduler config dict with the step-interval LambdaLR.

DetectionDataModule

Bases: LightningDataModule

Detection datamodule assembling the Phase 1 augmentation pipeline (WP-014).

Train and val splits are built from the COCO 2017 layout under data_root (train2017/ + annotations/instances_train2017.json and the val2017 equivalents); each split path may be overridden explicitly (used by the offline fixture tests, which have a single Roboflow-style split). A data_root written in the YOLO convention instead is read by :class:~lucid_yolo.data.yolo.YoloDetectionDataset — see the module docstring for the three-step dispatch (WP-099b). Training samples pass through :class:_TrainPipeline (mosaic/mixup/copy-paste/affine/hsv/flip + letterbox); validation is letterbox-only. No split is ever downloaded.

Parameters:

Name	Type	Description	Default
data_root	Path	Dataset root: a COCO 2017 root (holding train2017, val2017, annotations), any other layout :data:~lucid_yolo.data.layout.CANDIDATES names, or a YOLO root (a data.yaml beside images/labels trees).	required
batch_size	int	Samples per batch for both loaders.	required
num_workers	int | None	DataLoader worker processes. None (default) resolves to min(batch_size, cpu count), further capped so the in-flight worker queue fits the free /dev/shm budget on Linux (:func:_shm_capped_workers) — a many-core host with a large batch otherwise fills the shm tmpfs before the first step. Determinism of the seeded pipeline is guaranteed only at an explicit 0 (a single in-process generator).	required
variant	str	Model size letter selecting the augmentation strength policy ("n"…"x"); validated at construction.	required
img_size	int	Square letterbox side for every emitted sample. Defaults to 640, and must be a positive multiple of every head stride — validated at construction, for the reason variant is.	640
train_images_dir	Path | None	Override for the train images directory. Defaults to whichever of data_root/"train2017" and data_root/"train" :func:~lucid_yolo.data.layout.resolve_split finds, so a tiled oriented root needs no override at all.	None
train_ann_file	Path | None	Override for the train annotation JSON. Resolved alongside the images directory, never independently of it.	None
val_images_dir	Path | None	Override for the val images directory. Resolved as the train one is.	None
val_ann_file	Path | None	Override for the val annotation JSON. Resolved alongside the val images directory.	None
seed	int	Seed for the training pipeline's generator and loader shuffling.	0
pin_memory	bool | None	Whether loader batches land in page-locked host memory for async host-to-device copies. None (default) resolves to True exactly when CUDA is available — MPS and CPU runs get False (pinning buys nothing there and MPS warns on it).	None
prefetch_factor	int	Batches each worker keeps prefetched (defaults to :data:_PREFETCH_FACTOR); ignored at num_workers=0 where the DataLoader forbids it.	_PREFETCH_FACTOR
val_num_workers	int | None	Worker count for the validation loader only, capped by nothing — the flag an operator states to spend the whole machine on validation. None (default) resolves to an auto-chosen num_workers bounded by :func:_val_worker_target (WP-073, WP-102), and to a named num_workers bounded by the in-flight memory budget of :meth:_val_loader_workers (WP-103) — letterbox-only val samples need few workers, and both worker populations are resident at the epoch boundary where validation runs.	None
persistent_workers	bool	Keep loader workers alive across epochs. Defaults to False (WP-076): respawning costs seconds per epoch while a persistent pool accumulates per-worker memory for the whole run — the observed slow-creep OOM on long containerized runs.	False
cache_train_images	bool	Decode the training split once during :meth:setup and retain read-only RGB tensors — uint8 for 8-bit sources, uint16 (twice the bytes) for a 16-bit PNG. Defaults to False; enabling it trades the uncompressed split size in host RAM for removing image decode from every augmented draw in every epoch. :meth:setup only pays this cost for a "fit"/None stage, so a "validate"/"test"/"predict" call never decodes the train split it will not read. The decode draws a progress bar and reports the cache's size as a UserWarning — before decoding on a COCO root, from the annotation file's image sizes, and after it on a YOLO root, which states none — so an undersized host sees the figure rather than a silent stall ending in an out-of-memory kill.	False
cache_prepared_train_images	bool	Retain float32, fill-centred source images instead of repeatedly converting bytes and centring each random mosaic. Off by default. Takes precedence over cache_train_images, with a UserWarning at construction when both are set, and costs four bytes per pixel channel, per training record and training process; duplicate COCO records retain separate images, plus target metadata. YOLO labels are snapshotted with their images. The random geometry, fill-restored pixels, targets and RNG draws stay unchanged. Fit setup builds it once, drawing a progress bar; other stages release it. Host RAM must hold the split.	False
device_hsv	bool	Apply training HSV and quantization on CUDA, in batches, under float32 or mixed precision. Off by default; other backends and precisions retain worker-side colour augmentation. The workers keep their RNG draws and transport float32 images plus the stated gains, using four times the default image IPC and transfer bytes. The worker cap budgets float32 images only for a count chosen here (num_workers unset) in one process; a named num_workers is neither capped nor warned, and each DDP rank budgets the free shm on its own. Validation stays uint8.	False
fused_hsv	bool	Let the transfer hook restore a deferred-HSV batch with the fused Triton kernel where the host supports it (:mod:lucid_yolo.data._hsv_cuda). On by default, so a supported host takes it whenever Triton and a host C compiler are present; False keeps device_hsv on the native CUDA restore, to step around a toolchain or bisect a difference between the two. Pixels are identical either way. Has no effect without device_hsv.	True
rotated_targets	bool	Read every annotation's four-point ring as a rotated box (:class:~lucid_yolo.data.coco.CocoDetectionDataset's oriented mode, WP-088), so both loaders emit Targets whose rboxes pair 1:1 with their boxes. Off by default because it requires a quadrilateral ring per instance; the CLI turns it on for model.task="obb". Costs no extra transport — the packed batch has carried a rotated-box axis since WP-055.	False
mask_targets	bool	Rasterise the segmentation mask targets in the loader workers (:func:collate_detection) instead of leaving that CPU work in the training process. Off by default because it requires a polygon ring per box; the CLI turns it on for model.task="segment". Adds the mask stack to the batch transport — roughly 27 MB per batch of 32 at COCO instance density — so a container with a small /dev/shm may need fewer workers or a lower prefetch_factor. Unavailable on a YOLO root: that format carries no per-instance rings at all, so the request is refused as soon as the layout is known rather than left to produce a segmentation run supervised by nothing.	False
keypoint_targets	bool	Parse every annotation's COCO keypoints field (:class:~lucid_yolo.data.coco.CocoDetectionDataset's keypoints mode, WP-121), so both loaders emit Targets whose keypoints and keypoint_vis pair 1:1 with their boxes. Off by default because it requires a keypoints-annotated file; the CLI turns it on for model.task="keypoints". Adds a (sum_N, K, 2) coordinate axis and a (sum_N, K) visibility axis to the batch transport — small beside the images. Unavailable on a YOLO root, which carries no keypoint fields. Both loaders carry points the whole way (WP-132). The validation path is letterbox-only and needed nothing but :class:~lucid_yolo.data.letterbox.Letterbox; the training path also runs the mosaic, the fused affine, mixup and the horizontal flip, and each of those rebuilds its Targets, so each had to be taught the point channels explicitly. Unlike the letterbox they crop and filter, which is why they need A70's rule for a point warped off the canvas — it is carried through with its true coordinate rather than clamped or demoted, a case A66's v == 0 ("no annotation exists") does not answer. The flip needs more than a rule: it needs this dataset's own left/right pairing, which is why the reader publishes keypoint_flip_pairs (A64) rather than the transform assuming one.	False
layout	str | None	Which reader data_root calls for — "coco", "yolo", or None (the default) to infer it. Two things are unsayable without it, which is the whole of its justification: a root holding both conventions, which :func:~lucid_yolo.data.layout.detect_layout refuses to break a tie on because the two readings are different label spaces; and a YOLO root whose data.yaml points its splits outside the convention table, which the reader resolves and the probe cannot see. Stating "coco" or "yolo" skips the probe entirely, so the reader — not this dispatch — names any path it then fails to open.	None
epoch_prefetch	bool	Queue up to num_workers x prefetch_factor batches of the next training epoch during the last batch and validation, using the existing persistent workers. Opt-in: queued training inputs can remain resident during validation. Only full, single-CUDA epochs with the module-owned sampler use the early reset; limits, step bounds, overfit, reloads and distributed sampling stay native. CUDA is the only accelerator it was measured on (an L4); others are unmeasured, not known to fail. A request the attached trainer cannot honour warns once, on rank zero, naming the first unmet condition.	False

Examples:

>>> DetectionDataModule  # doctest: +SKIP
>>> # dm = DetectionDataModule(root, batch_size=2, num_workers=0, variant="n")
>>> # dm.setup("fit"); images, targets = next(iter(dm.train_dataloader()))

layout property

layout: DatasetLayout

Return the layout this datamodule reads, probing the root once if need be.

Returns:

Type	Description
DatasetLayout	The resolved :class:~lucid_yolo.data.layout.DatasetLayout, cached after the
DatasetLayout	first call so the probe's verdict cannot change mid-run.

Raises:

Type	Description
FileNotFoundError	If the root satisfies neither convention; the message names both and every path tried (:func:~lucid_yolo.data.layout.detect_layout).
ValueError	If it satisfies both, or asks for something that layout cannot serve.

Examples:

>>> DetectionDataModule.layout  # doctest: +SKIP
>>> # DetectionDataModule(yolo_root, batch_size=2, num_workers=0, variant="n").layout
>>> # <DatasetLayout.YOLO: 'yolo'>

mosaic_prob property

mosaic_prob: float

Return the training pipeline's current mosaic probability.

Returns:

Type	Description
float	The probability with which each training sample is built from a
float	four-image mosaic (1.0 until the close_mosaic schedule flips it).

Raises:

Type	Description
RuntimeError	If :meth:setup has not built the training pipeline yet.

Examples:

>>> DetectionDataModule.mosaic_prob  # doctest: +SKIP
>>> # dm.setup("fit"); dm.mosaic_prob
>>> # 1.0

prepare_data

prepare_data() -> None

No-op: datasets are provisioned out of band, never downloaded (AGENTS.md section 3).

setup

setup(stage: str | None = None) -> None

Build the train pipeline and the letterbox-only val dataset.

The layout is resolved here (:attr:layout) rather than at construction, because this is the first moment the root is required to exist: a root that satisfies neither convention raises before a single image is decoded, naming both conventions and every path tried.

Parameters:

Name	Type	Description	Default
stage	str | None	The Lightning stage ("fit"/"validate"/…). Both splits are always built, so a repeated call yields the same splits and a fresh train pipeline — mosaic probability and draw generator back at their initial state. cache_train_images only decodes the train split when stage is "fit" or None: a "validate"/"test"/"predict" call never pays to cache a split it will not read, and releases a cache an earlier fit-stage call built. A repeat "fit"/None call with the cache on keeps the already-decoded train reader rather than decoding the split again beside it, so the result is idempotent and the decode cost is paid once, not once per call. A call that raises installs nothing it built: the val dataset and both caches keep their previous values, and no train pipeline is left behind, so :meth:train_dataloader raises rather than serving the pipeline an earlier stage built.	None

Raises:

Type	Description
FileNotFoundError	If the root satisfies neither dataset convention.
ValueError	If it satisfies both, or the resolved layout cannot serve a requested modality.

on_after_batch_transfer

on_after_batch_transfer(batch: tuple[Tensor, PackedTargets] | tuple[Tensor, list[Targets]], dataloader_idx: int) -> tuple[Tensor, list[Targets], list[Tensor] | None]

Restore the float images and list[Targets] consumer contract on-device.

Lightning calls this after moving the batch to the accelerator, so both the uint8 images and the :class:PackedTargets tensor fields are already on-device: the images are dequantized to the attached module's dtype (:meth:_consumer_dtype — code / 255 at that precision, :func:_dequantize_images) and the packed targets are unpacked, all on the same device the module trains on, with no extra host-to-device hop. A deferred training batch first applies its stated HSV gains and quantizes the result, then reaches this same dequantization. Matching the module dtype keeps a "16-true"/"bf16-true" run's half weights and inputs aligned; "16-mixed" leaves the module float32 and lets autocast cast. Already-float images and already-list targets (a non-packed direct feed) pass through untouched, so the hook is safe to run over either form. Both the train and val loaders share this single restore. On CUDA the restored images are made channels_last, the layout the module's parameters take in :meth:DetectionLitModule.setup; masks and every other device keep the default layout.

Parameters:

Name	Type	Description	Default
batch	tuple[Tensor, PackedTargets] | tuple[Tensor, list[Targets]]	The transferred (uint8 images, PackedTargets) transport batch, or an already-restored (float images, list[Targets]) batch.	required
dataloader_idx	int	The loader index (unused; the conversion is uniform).	required

Returns:

Type	Description
Tensor	The (images, list[Targets], masks) consumer batch — [0, 1] images at the module's dtype, the
list[Targets]	ragged per-image targets, and the per-image mask stacks a segmentation loader rasterised in its workers
list[Tensor] | None	(None on a detection loader, or on any already-restored feed, which is the signal the module
tuple[Tensor, list[Targets], list[Tensor] | None]	rasterises the masks itself).

Raises:

Type	Description
TypeError	If a deferred-HSV batch arrives cast away from float32 images and float64 gains (:func:_restore_images).

Examples:

>>> DetectionDataModule.on_after_batch_transfer  # doctest: +SKIP
>>> # images, targets = dm.on_after_batch_transfer(next(iter(dm.val_dataloader())), 0)

train_dataloader

train_dataloader() -> DataLoader[tuple[Tensor, PackedTargets]]

Return the shuffled training loader with the packed batch transport.

DataLoader's type parameter describes dataset items, while this collate returns a different form. The cast records the collated boundary rather than exposing the pipeline's private deferred-sample union to consumers. The loader collates through :func:_collate_training, the private twin of :func:collate_detection that also accepts that union, so the public collate keeps its (image, Targets) signature.

val_dataloader

val_dataloader() -> DataLoader[tuple[Tensor, PackedTargets]]

Return the unshuffled validation loader (letterbox-only samples).

The loader runs on its own worker count — val_num_workers, defaulting to a pixel-scaled cap on an auto-chosen train count (WP-073, WP-102) and to the memory budget of :meth:_val_loader_workers on a named one (WP-103): letterbox only samples need a fraction of the train pipeline's CPU, and both worker populations are resident at the epoch boundary where validation runs.

Annotated with the collated (images, PackedTargets) transport, as :meth:train_dataloader is: the dataset yields (image, Targets) pairs, and the cast records the form :func:collate_detection actually hands the trainer.

set_mosaic_prob

set_mosaic_prob(prob: float) -> None

Set the training pipeline's mosaic probability.

This is the seam CloseMosaicCallback (WP-036) drives to disable mosaic for the final training epochs: it sets prob=0. The value lands on the main-process pipeline, which loader worker processes copy when they start: in-process loading (num_workers=0) uses it from the next sample, and a pool started after this call — every later epoch's pool under the default persistent_workers=False — samples with it. Persistent workers keep the copy they started with and never see the change; with epoch_prefetch also requested, a change of value warns once, since the next epoch's first batches are also already queued by then. Repeat calls are harmless.

Parameters:

Name	Type	Description	Default
prob	float	The new mosaic probability, forwarded to :attr:_TrainPipeline.mosaic_p.	required

Raises:

Type	Description
RuntimeError	If :meth:setup has not built the training pipeline yet.

Examples:

>>> DetectionDataModule.set_mosaic_prob  # doctest: +SKIP
>>> # dm.setup("fit"); dm.set_mosaic_prob(0.0); dm.mosaic_prob
>>> # 0.0

Callbacks

EMACallback

Bases: Callback

Track an exponential-moving-average shadow of the weights and eval with it.

A shadow copy of the module's floating-point parameters and buffers (batch-norm running statistics included; integer buffers such as num_batches_tracked are left out — they carry no meaningful average) is created on :meth:on_fit_start on the module's device. After each optimizer step (:meth:on_train_batch_end, honoring update_every) the shadow is blended toward the live weights under :func:torch.no_grad with a warmup-ramped decay

d_t = decay * (1 - exp(-t / tau)) shadow <- d_t * shadow + (1 - d_t) * live

where t is the number of updates performed so far. Tensors sharing a device and dtype are submitted to one foreach interpolation: the same recurrence with fewer accelerator launches, but not bit-identical to the per-tensor form in any dtype. It agrees within the test tolerance for fp32 and fp64, and a half-precision shadow rounds once per step instead of twice (a difference of about 1 ulp). A tensor whose shadow and live dtypes differ keeps the per-tensor recurrence. The ramp starts the effective decay near zero so the early shadow tracks the fast-moving weights closely, then relaxes toward the nominal decay as training settles (blueprint D4). Around the validation loop the callback swaps the shadow in (:meth:on_validation_start) and the live weights back out (:meth:on_validation_end), so validation metrics are EMA metrics — the "used for eval" contract of D4. The shadow and update counters are checkpointed via :meth:state_dict/:meth:load_state_dict, so a save/resume round trip continues the same average rather than restarting it.

Cadence is optimizer steps, not batches. update_every and tau are both counted in optimizer steps, so accumulate_grad_batches > 1 does not multiply the number of blends: :meth:on_train_batch_end blends only on a batch that actually advanced trainer.global_step, which under accumulation is the last batch of each accumulation window. That hook rather than on_before_zero_grad because Lightning calls the latter at the start of an accumulation window, before that window's training_step — it would blend pre-step weights and never see the state left by the final optimizer step of a run.

Callback ordering. The weight swap is deliberately hook-scoped, not run-scoped, so it does not depend on how Lightning happens to order the callback list: the shadow is installed in :meth:on_validation_start and taken out again in :meth:on_validation_end (and in :meth:on_exception, so a failed validation cannot leave EMA weights wearing the live model's identity). Checkpoints written by :class:~pytorch_lightning.callbacks.ModelCheckpoint therefore hold the raw trained weights in state_dict, with the EMA available separately under this callback's state_dict — the split :func:~lucid_yolo.eval.checkpoint.load_eval_module relies on when it overlays the shadow on request. ModelCheckpoint saves from on_validation_end too, and Lightning moves it to the end of the callback list, so within a single validation the restore runs first; the two facts are recorded together because only the second is Lightning's to change.

The exact decay/tau constants are unpublished in the source papers; the defaults are the conventional YOLO-lineage values and are documented here as a small, self-contained assumption (D4 names EMA explicitly, so no ASSUMPTIONS.md row is warranted).

Parameters:

Name	Type	Description	Default
decay	float	Nominal per-update decay the warmup ramp relaxes toward. Must be in [0, 1). Defaults to 0.9999.	_DEFAULT_DECAY
tau	int	Warmup time constant (in updates, and so in optimizer steps) of the decay ramp; larger values ramp in more slowly. Must be positive. Defaults to 2000.	_DEFAULT_TAU
update_every	int	Refresh the shadow every update_every optimizer steps — not every update_every batches. Must be at least 1. Defaults to 1.	_DEFAULT_UPDATE_EVERY

Raises:

Type	Description
ValueError	If decay is outside [0, 1), tau is not positive, or update_every is less than 1.

Examples:

>>> from lucid_yolo.ptl.callbacks import EMACallback
>>> callback = EMACallback(decay=0.99, tau=100)
>>> round(callback._decay_at(1), 6)  # ramped far below the nominal decay
0.009851
>>> callback._decay_at(1) < callback.decay
True

on_fit_start

on_fit_start(trainer: Trainer, pl_module: LightningModule) -> None

Create (or, on resume, relocate) the shadow on the module's device.

A shadow restored from a checkpoint via :meth:load_state_dict is kept — only its device is refreshed — so resuming continues the saved average; otherwise a fresh shadow is cloned from the current floating-point parameters and buffers.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer (unused; part of the hook signature).	required
pl_module	LightningModule	The module whose weights are shadowed.	required

on_train_start

on_train_start(trainer: Trainer, pl_module: LightningModule) -> None

Seed the optimizer-step watermark from the trainer, so a resume does not double-count.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer, read for its global_step.	required
pl_module	LightningModule	The module being trained (unused; part of the hook signature).	required

on_train_batch_end

on_train_batch_end(trainer: Trainer, pl_module: LightningModule, outputs: Any, batch: Any, batch_idx: int) -> None

Blend the shadow toward the live weights once per optimizer step.

A batch that did not advance trainer.global_step — every batch of an accumulate_grad_batches window but its last — leaves the shadow untouched, so update_every and the tau ramp are counted in optimizer steps as documented rather than in batches. By this hook the step has already been applied, so the blend reads post-step weights.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer, read for its global_step.	required
pl_module	LightningModule	The module whose weights are shadowed.	required
outputs	Any	The step outputs (unused; part of the hook signature).	required
batch	Any	The batch just processed (unused; part of the hook signature).	required
batch_idx	int	Index of the batch within the epoch (unused).	required

on_validation_start

on_validation_start(trainer: Trainer, pl_module: LightningModule) -> None

Stash the live weights and load the shadow so validation runs on the EMA.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer (unused; part of the hook signature).	required
pl_module	LightningModule	The module to load the shadow into.	required

on_validation_end

on_validation_end(trainer: Trainer, pl_module: LightningModule) -> None

Restore the live weights stashed in :meth:on_validation_start.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer (unused; part of the hook signature).	required
pl_module	LightningModule	The module to restore the live weights into.	required

on_exception

on_exception(trainer: Trainer, pl_module: LightningModule, exception: BaseException) -> None

Take the shadow back out of the module when a run dies mid-validation.

Without this the swap of :meth:on_validation_start is undone only by the normal completion path, so an exception raised inside the validation loop propagates with the EMA weights still installed under the raw model's identity — anything the failure handler then inspects, saves or resumes from would be the shadow rather than the trained weights. Restoration is idempotent, so the ordinary path running first (or this hook firing outside validation) is harmless.

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer (unused; part of the hook signature).	required
pl_module	LightningModule	The module to restore the live weights into.	required
exception	BaseException	The exception being propagated (unused; part of the hook signature).	required

state_dict

state_dict() -> dict[str, Any]

Return the shadow and update counters for Lightning to checkpoint.

Returns:

Type	Description
dict[str, Any]	A dict with the shadow tensor mapping (or None before on_fit_start)
dict[str, Any]	and the num_updates / steps_seen counters. The shadow tensors are
dict[str, Any]	cloned so a checkpoint captures the snapshot at call time, not a later state.

load_state_dict

load_state_dict(state_dict: dict[str, Any]) -> None

Restore the shadow and update counters from a checkpoint.

A checkpoint written before the cadence was fixed to optimizer steps carries the counter under batches_seen; it is read as the update_every phase it in fact was, so resuming such a run continues rather than restarts. The two counts differ only when that run used accumulate_grad_batches > 1.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	A mapping produced by :meth:state_dict.	required

CloseMosaicCallback

Bases: Callback

Disable mosaic augmentation for the final close_mosaic training epochs.

Mosaic (four-image compositing) runs at p=1.0 for most of training but is switched off for the last close_mosaic epochs so the network finishes on natural, un-composited images (blueprint sec. 5.9; [R1] Tables S3/S6). At the start of every training epoch this callback checks whether the current epoch has reached the close-mosaic window (current_epoch >= max_epochs - close_mosaic) and, if so, sets the training pipeline's mosaic probability to 0 via the datamodule's :meth:~lucid_yolo.ptl.datamodule.DetectionDataModule.set_mosaic_prob seam. The flip is idempotent — re-setting 0 each subsequent epoch is harmless — and is logged exactly once, on the epoch it first takes effect (rank zero only).

Parameters:

Name	Type	Description	Default
close_mosaic	int	Number of final epochs to train with mosaic disabled. 0 never disables mosaic; a value >= max_epochs disables it from the first epoch. Must be non-negative.	_DEFAULT_CLOSE_MOSAIC

Raises:

Type	Description
ValueError	If close_mosaic is negative.

Examples:

>>> from lucid_yolo.ptl.callbacks import CloseMosaicCallback
>>> callback = CloseMosaicCallback(close_mosaic=10)
>>> callback.close_mosaic
10

on_train_epoch_start

on_train_epoch_start(trainer: Trainer, pl_module: LightningModule) -> None

Disable mosaic once the close-mosaic window is reached (idempotent).

Parameters:

Name	Type	Description	Default
trainer	Trainer	The active trainer, read for current_epoch, max_epochs, and the attached datamodule.	required
pl_module	LightningModule	The module being trained (unused; part of the hook signature).	required

Detection loss

DualBranchLoss takes one BranchPrediction(logits, boxes) per branch, followed by shared anchor points and padded targets. Its output records the attribute alpha float by default, or the supplied alpha tensor when overridden.

BranchPrediction

Bases: NamedTuple

Pair one detection branch's raw class logits and decoded boxes.

Attributes:

Name	Type	Description
logits	Tensor	Raw class logits with shape (B, A, C).
boxes	Tensor	Decoded xyxy boxes with shape (B, A, 4).

Examples:

>>> import torch
>>> prediction = BranchPrediction(torch.zeros(1, 4, 2), torch.zeros(1, 4, 4))
>>> prediction.logits.shape, prediction.boxes.shape
(torch.Size([1, 4, 2]), torch.Size([1, 4, 4]))

DualBranchLoss

Compose the one-to-many and one-to-one detection-branch losses.

Holds two :class:~lucid_yolo.losses.detection_loss.DetectionBranchLoss instances (identical gains), a dense o2m assigner (:class:~lucid_yolo.assign.stal.SmallTargetAssigner, topk = 10), and a unique o2o assigner (:class:~lucid_yolo.assign.one_to_one.UniqueAssigner, topk = 7 reduced to one). The object is stateless apart from its gains, assigners, and the settable :attr:alpha, so a single instance is reused across the training run. It is coordinate-scale agnostic in the same way as the per-branch loss: predicted and ground-truth boxes only need to share a frame (see :class:~lucid_yolo.losses.detection_loss.DetectionBranchLoss).

Parameters:

Name	Type	Description	Default
box_gain	float	Weight on each branch's Complete-IoU term (A13 default 7.5).	7.5
cls_gain	float	Weight on each branch's classification term (A13 default 0.5).	0.5
l1_gain	float	Weight on each branch's L1 box term (A13 default 6.0).	6.0
o2m_topk	int	Candidate count for the dense o2m assignment (default 10).	10
o2o_topk	int	Candidate-set size for the o2o assignment before the unique topk2 = 1 reduction (default 7).	7
alpha	float	Initial branch weight for L = alpha * L_o2m + (1 - alpha) * L_o2o (static default 0.5; the WP-035 schedule sets it per epoch).	0.5
s_min	float	Small-target surrogate threshold shared by both assigners (the smallest stride, 8.0 at 640 input).	8.0
s_ref	float	Small-target surrogate replacement size shared by both assigners (the next stride, 16.0 at 640 input).	16.0
fused_assignment	bool	Request the optional CUDA assignment passes on both branches. Default False keeps the existing schedule.	False

Examples:

>>> import torch
>>> from lucid_yolo.assign import make_anchor_points
>>> loss = DualBranchLoss()
>>> points, _ = make_anchor_points([(2, 2)], [8])  # 4 anchors
>>> logits = torch.zeros(1, 4, 1)
>>> boxes = torch.zeros(1, 4, 4)
>>> empty_boxes = torch.zeros(1, 0, 4)
>>> empty_labels = torch.zeros(1, 0, dtype=torch.long)
>>> empty_mask = torch.zeros(1, 0, dtype=torch.bool)
>>> out = loss(
...     BranchPrediction(logits, boxes), BranchPrediction(logits, boxes),
...     points, empty_boxes, empty_labels, empty_mask
... )
>>> bool(torch.isfinite(out.total))  # no ground truth: finite, cls-only
True

__call__

__call__(o2m: BranchPrediction, o2o: BranchPrediction, anchor_points: Tensor, gt_boxes: Tensor, gt_labels: Tensor, gt_mask: Tensor, strides: Tensor | None = None, gt_rboxes: Tensor | None = None, *, alpha: Tensor | None = None) -> DualLossOutput

Score both branches against their own assignments and combine them.

Parameters:

Name	Type	Description	Default
o2m	BranchPrediction	Raw logits and decoded boxes of the one-to-many branch.	required
o2o	BranchPrediction	Raw logits and decoded boxes of the one-to-one branch.	required
anchor_points	Tensor	(A, 2) anchor-centre (x, y) pixels shared by both branches (see :func:lucid_yolo.assign.make_anchor_points).	required
gt_boxes	Tensor	(B, N, 4) padded ground-truth boxes in xyxy pixels.	required
gt_labels	Tensor	(B, N) long class ids; padded slots are ignored.	required
gt_mask	Tensor	(B, N) bool marking real ground truths.	required
strides	Tensor | None	Optional (A,) per-anchor level strides, forwarded to both branch losses so their L1 terms are measured in stride units (A13 revision, WP-078). The training path always passes strides.	None
gt_rboxes	Tensor | None	Optional (B, N, 5) rotated ground truths in the long-edge convention, entry (b, n) describing the same object as gt_boxes[b, n]. Forwarded verbatim to both assigners, where it switches candidacy to point-in-rotated-rect containment (WP-061, A25) and nothing else: the IoU, the alignment metric, the targets and the weights keep running on the axis-aligned envelopes. Omitted (the default) leaves the assignment, and therefore every tensor operation on this path, exactly as it was — which is what keeps the detection objective bit-identical.	None
alpha	Tensor | None	Optional 0-d float64 tensor carrying the branch weight for this call in place of :attr:alpha. The compiled detection step (WP-184) passes it because a Python float read inside the traced region is a Dynamo guard on its value, and the WP-035 schedule changes it every epoch — one recompile per epoch until cache_size_limit dropped the region to eager (WP-187). The blend casts it to each branch total's dtype after forming 1 - alpha in double, which is exactly the arithmetic the Python-float form does, so the total is bit-identical either way. Omitted (the default) blends with the attribute.	None

Returns:

Name	Type	Description
A	DualLossOutput	class:DualLossOutput with the combined total, each branch's
	DualLossOutput	class:~lucid_yolo.losses.detection_loss.DetectionLossOutput, the alpha used, and both branches'
	DualLossOutput	assignments (so a mask term can supervise exactly these positives instead of assigning again).

Examples:

>>> import torch
>>> from lucid_yolo.assign import make_anchor_points
>>> loss = DualBranchLoss()
>>> loss.alpha = 0.25  # e.g. a scheduled value from WP-035
>>> points, _ = make_anchor_points([(2, 2)], [8])
>>> logits = torch.zeros(1, 4, 1, requires_grad=True)
>>> boxes = torch.zeros(1, 4, 4, requires_grad=True)
>>> gt = torch.tensor([[[0.0, 0.0, 8.0, 8.0]]])
>>> out = loss(
...     BranchPrediction(logits, boxes), BranchPrediction(logits, boxes),
...     points, gt, torch.tensor([[0]]), torch.tensor([[True]])
... )
>>> out.total.backward()
>>> bool(torch.isfinite(logits.grad).all())
True

DualLossOutput dataclass

Combined dual-branch loss, its two per-branch breakdowns, and both assignments.

total is the alpha-weighted combination of the two branch totals; the o2m and o2o fields carry each branch's full :class:~lucid_yolo.losses.detection_loss.DetectionLossOutput (its pre-gain box/cls/l1 terms and gain-weighted total), and alpha is the branch weight used for this call.

The two :class:~lucid_yolo.assign.tal.AssignResult values are returned rather than kept private because any other term supervising the same positives — the WP-087 instance-mask loss is the first — must score the ground truth each anchor was actually assigned to here. Recomputing the assignment at the call site would create a second selection path free to disagree with this one, and a mask scored against the wrong (but plausible) instance is invisible in the loss value.

Attributes:

Name	Type	Description
total	Tensor	alpha * o2m.total + (1 - alpha) * o2o.total; the scalar to backpropagate.
o2m	DetectionLossOutput	One-to-many branch loss output (dense topk = 10 assignment).
o2o	DetectionLossOutput	One-to-one branch loss output (unique topk = 7 -> 1 assignment).
alpha	float | Tensor	Branch weight applied to the o2m total (o2o gets 1 - alpha) — the :attr:DualBranchLoss.alpha float when no override is supplied, otherwise the override tensor itself, without a read back to the host.
o2m_assign	AssignResult	The dense assignment the o2m branch was scored against.
o2o_assign	AssignResult	The unique assignment the o2o branch was scored against.