🏋️ 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
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

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 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
float	underlying loss attribute so a scheduler 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

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.

val/mAP and val/segm_mAP come from :mod:torchmetrics metrics, which synchronise across ranks themselves. The oriented and keypoint protocols do not: :attr:_val_rotated_preds and :attr:_val_keypoint_preds are ordinary Python lists, so on more than one process each rank would score and log its own shard as if it were the whole split. That is not recoverable after the fact — average precision ranks every detection of the split against every other by confidence, so averaging per-rank AP values is a different quantity, not an approximation of the right one. Rather than log a plausible wrong number, these tasks refuse to start distributed; "detect" and "segment" are unaffected and stay free to scale.

Lifting this needs the buffers gathered with image identity (so sampler padding's duplicate images are dropped rather than double-counted) and scored once on one rank — the work is real, so the limit is enforced here instead of assumed.

Parameters:

Name	Type	Description	Default
stage	str	The Trainer stage being set up ("fit", "validate", "test" or "predict"). "predict" runs no validation protocol and is exempt.	required

Raises:

Type	Description
RuntimeError	If the task is "obb" or "keypoints" and the trainer spans more than one process.

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

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
	DualHeadOutput	predictions (raw class logits and raw ltrb distances) of both
	DualHeadOutput	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
	SegmentOutput	head's dense predictions (including both branches' mask
	SegmentOutput	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
BranchOutput	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.

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(),
	OptimizerLRScheduler	alone or inside a Lightning optimizer/scheduler config dict with the
	OptimizerLRScheduler	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
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
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

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.	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
list[Targets]	at the module's dtype, the ragged per-image targets, and the per-image
list[Tensor] | None	mask stacks a segmentation loader rasterised in its workers (None on
tuple[Tensor, list[Targets], list[Tensor] | None]	a detection loader, or on any already-restored feed, which is the signal
tuple[Tensor, list[Targets], list[Tensor] | None]	the module 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 change is consulted on the next sample the pipeline assembles. 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_checkpoint_model 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