diff --git a/CHANGELOG.md b/CHANGELOG.md index 8ea9e46..8952d47 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ All notable changes to lucid-yolo are documented here, following the Keep a Chan ### Added +- `capture_network` now replays shared backbone and neck stages for segmentation, independently of `capture_segment_heads`. The default remains off; modified heads/forwards, accumulation and unsupported signatures retain native execution, with transient graph state cleared at fit boundaries (WP-193g). + - **Opt-in CUDA network replay for oriented training.** `--model.capture_network true` records the existing backbone, neck and both oriented head branches, including their angle outputs. The native objective and state updates remain outside capture. Accumulation and overridden objectives keep the native route; the default stays off. The speed claim is limited to a small-batch workload (Nano OBB at batch 4, BF16, one L4): a single batch-8 pair changed fit time by only 4.10%, and replay raised peak reserved memory from 3974 to 4300 MiB, so check throughput and memory on your own batch size. See [WP-200](docs/work_packages/wp-200.md) for complete-fit measurements and limits. - **Opt-in CUDA replay of detection assignment and loss.** `--model.capture_loss true` records the existing loss forward/backward with a masked target workspace that grows from observed counts. Current label checks and schedules stay live; returned values own their storage and fit exit releases the recording. Default off, compatible with network and orthogonalization replay. See [WP-199](docs/work_packages/wp-199.md) for measurements and limits. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index eba1062..6a6b89e 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -357,6 +357,7 @@ Opened by issue 2 on the tracker: R1 sec. 3.4.2 publishes the pose objective thi | 193d | `perf(ptl): replay shared features in small-batch CUDA pose training` [DATA][GPU] | Pose replays the existing backbone and neck while its head stays native, preserving absent point gradients. The default stays off; accumulated runs, a fully frozen backbone and neck, and heads with changed topology, hooks or forwards keep the native route · [log](work_packages/wp-193d.md#wp-193d) | `tests/ptl/test_pose_network_capture.py`, whose synthetic `layout="empty"` case of `test_pose_trainer_matches_native` evidences the absent point-stem gradients of a batch with no assigned positives — the public-COCO parity run does not; public-COCO byte parity and five isolated-L4 complete fits with 11.17% median time reduction at batch 4 (one exploratory batch-16 pair: 2.50%; peak allocated memory down, reserved up); torch 2.4/2.9 focused checks and full pre-commit; full `make gate`, accelerator goldens and final accuracy not run | 193c | 🔄 | | 193e | `perf(ptl): replay independent segmentation predictions` [DATA][GPU] | Opt-in BF16 CUDA prediction groups preserve native gradient absence and per-stem accumulation. One topology pass and an attribute certificate retain live storage and observer checks; default, evaluation and unsupported modes stay native · [log](work_packages/wp-193e.md#wp-193e) | `tests/ptl/test_segment_head_capture.py`; exact public-COCO states; five paired isolated-L4 complete fits; full pre-commit; full `make gate` not run | 193c | 🔄 | | 193f | `perf(ptl): replay independent CUDA pose prediction stems` [DATA][GPU] | The existing pose flag records eighteen native stems independently through the shared prediction runner, preserving absent point gradients and bounded fit ownership · [log](work_packages/wp-193f.md#wp-193f) | Focused CPU/CUDA, real replay and lifecycle regressions; five isolated-L4 complete pairs with 12.97% median time reduction at batch 4 and byte-identical states; reserved memory increases; full gate, aggregate coverage and goldens not run | 193d,193e | 🔄 | +| 193g | `perf(ptl): replay shared segmentation features` [DATA][GPU] | Opt-in backbone/neck replay preserves independent segmentation prediction gradients. Combines with head replay; accumulation and modified forwards stay native, with one fit-scoped signature · [log](work_packages/wp-193g.md#wp-193g) | `tests/ptl/test_segment_network_capture.py`; public-COCO byte parity; five paired isolated-L4 complete fits; focused tooling; full `make gate` and accelerator goldens not run (offline CI green at PR head 8ebf10e) | 193d, 193e | 🔄 | | 194 | `refine(tests): static checks to hooks, plain test names` | File-content tests leave pytest: the `_viz`/matplotlib import boundary becomes ruff `TID251`, the top-k and packaged-recipe text pins the `content-pins` hook, and live-tree "hook passes" copies are dropped. Test names lose `a an the its their is are does that` · [log](work_packages/wp-194.md#wp-194) | `scripts/_tests/test_audit_content_pins.py`; collection count unchanged; `make gate` green | 193 | ✅ | | 195 | `perf(data): batch rotated-box clipping and refitting` [DATA][GPU] | Finite float32 CPU augmentation boxes without gradients share the existing clip and orientation-preserving fit in equal-vertex-count groups. Ring order and instance order are preserved; other inputs keep the scalar route · [log](work_packages/wp-195.md#wp-195) | `tests/data/test_rotated_clip_batch.py`; unchanged public-DOTA training states and five paired isolated-L4 complete fits; focused data tests and full pre-commit; full `make gate` not run | 194 | 🔄 | | 196 | `docs(site): tabs, method pages, one file per WP` | The single left menu listed every register at once, three pages ran past 500 lines, the site never explained the method, and findings and results were copied between pages. A four-tab top bar (Home, Results, Models, Project); two Method pages with animated SVGs (how the detector works; NMS against NMS-free). The report splits into one page per section; each WP's findings move to `work_packages/wp-.md`, which the research and engineering phase pages include by section (D28), and the Project tab lists phases first. Model cards share one layout with an animated SVG, and include the report's result tables rather than copying them. Every WP named in the docs links to its page, each WP page opens with a roadmap-derived header (logs, dependencies, neighbours), and the README's Mermaid diagram becomes the animated SVG. For search engines and agents, a build hook gives every page a distinct title and description, Open Graph tags with a static preview card, JSON-LD on the home page, a markdown copy at `/index.md`, and `llms.txt` / `llms-full.txt`; the README links the site and `CITATION.cff` lands · [decision](DECISIONS.md) | `make docs` strict-green (every anchor and include resolves); `audit_wp_links.py` (`wp-links`) refuses an unlinked mention or a stale WP header; `audit_docs_site.py` honours `not_in_nav`; `audit_docs_present.py` checks each roadmap log link against its own target file and `check_work_package_files` refuses unnumbered or unincluded WP files; `audit_figure_captions.py` reads the report pages; `_WP_FLOOR` 196, `_DECISION_FLOOR` 28; `make gate` green | 195 | ✅ | diff --git a/docs/TRAINING.md b/docs/TRAINING.md index ff6de8c..6996dbc 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -69,7 +69,7 @@ lucid-eval --checkpoint --data_root /content/coco2017 --output seg_ ### Independent segmentation predictions on CUDA -`--model.capture_segment_heads true --trainer.precision bf16-mixed` records the existing segmentation prediction routines, including their first-order backward work. The default is `false`. Box/class stems remain independent; coefficients, prototypes and semantic logits have separate backward nodes so an empty mask objective keeps absent gradients. Common features, prototype fusion, assignment, losses, state updates and schedules stay native. +`--model.capture_segment_heads true --trainer.precision bf16-mixed` records the existing segmentation prediction routines, including their first-order backward work. The default is `false`. Box/class stems remain independent; coefficients, prototypes and semantic logits have separate backward nodes so an empty mask objective keeps absent gradients. With this option alone, common features, prototype fusion, assignment, losses, state updates and schedules stay native. Enable `capture_network` separately to replay the common features too. Byte-exact agreement with native training is verified only on NVIDIA L4 (sm_89) with torch 2.9.1+cu129; other GPUs and torch builds are unverified. The route requires PyTorch 2.9 or newer, one CUDA device, automatic optimisation and `accumulate_grad_batches=1`. Gradient accumulation stays native because replay was validated only at `accumulate_grad_batches=1`, not because of a known defect in the recorded backward. CPU, MPS, FP32, FP16 and older BF16 runtimes warn and use native calls; evaluation stays native too. Custom or changed prediction dispatch, hooks, frozen prediction parameters, aliased storage and retained feature gradients retire replay for the fit. Each prediction group records one input signature per fit (shape, stride, dtype, device, gradient flag and CUDA stream); a call with any other signature, a partial batch included, runs the native head without adding another graph, and the recorded graph pools stay reserved until the fit ends, so a run whose batch shapes keep changing pays that memory without replaying. Copying, checkpoints and fit teardown retain the original parameter names and no recording. @@ -145,10 +145,12 @@ For `task: keypoints`, the same flag records backbone and neck plus eighteen ind For `task: obb`, the flag records backbone, neck and the complete dual head, including both angle outputs. The native classification and oriented losses stay outside capture. Their empty-target expressions retain zero angle gradients. Overridden training consumers, classification losses or oriented helpers keep the complete native call; gradient accumulation also stays native. The oriented speed-up is likewise limited to a small-batch workload (Nano OBB at batch 4, BF16, one L4): a single batch-8 pair changed fit time by only 4.10%, and replay raised peak reserved memory from 3974 to 4300 MiB, so check throughput and memory on your own batch size. The measurements and remaining limits are in [WP-200](work_packages/wp-200.md). -This path requires single-device CUDA automatic optimisation and `task: detect`, `task: keypoints` or `task: obb`. BF16 replay requires a final PyTorch 2.9 or newer release; 2.9 pre-release and NGC `a0` builds are excluded by design. 2.9 is the first verified runtime: older BF16 runs warn and stay eager because the 2.4 comparison did not meet the exact eager oracle. FP32 and FP16 replay of detection, oriented detection and pose shared features are checked on 2.4 as well; pose stem replay has run on 2.9.1 only ([WP-193f](work_packages/wp-193f.md)). CPU, MPS and other strategies warn and stay eager; other tasks or `compile_step` raise at construction. One image signature is recorded per fit. A partial batch uses eager execution without growing a graph cache. Changes to the recorded stages’ parameter/buffer storage, module modes, hooks or precision retire replay for the rest of that fit. The supplied network must keep its ordinary stage implementations and static Python configuration. Root forward overrides, stage instance-forward overrides and registered hooks use the eager path. Anomaly detection (`--trainer.detect_anomaly true`) also stays eager so backward NaN checks run normally. A failure while recording raises instead of falling back to eager (RNG state and buffers are restored first), so disable `--model.capture_network` if recording fails on your model. Higher-order differentiation is not supported. +This path requires single-device CUDA automatic optimisation and `task: detect`, `task: keypoints`, `task: obb` or `task: segment`. BF16 replay requires a final PyTorch 2.9 or newer release; 2.9 pre-release and NGC `a0` builds are excluded by design. 2.9 is the first verified runtime: older BF16 runs warn and stay eager because the 2.4 comparison did not meet the exact eager oracle. FP32 and FP16 replay of detection, oriented detection and pose shared features are checked on 2.4 as well; pose stem replay has run on 2.9.1 only ([WP-193f](work_packages/wp-193f.md)). CPU, MPS and other strategies warn and stay eager; other tasks or `compile_step` raise at construction. One image signature is recorded per fit. A partial batch uses eager execution without growing a graph cache. Changes to the recorded stages’ parameter/buffer storage, module modes, hooks or precision retire replay for the rest of that fit. The supplied network must keep its ordinary stage implementations and static Python configuration. Root forward overrides, stage instance-forward overrides and registered hooks use the eager path. Anomaly detection (`--trainer.detect_anomaly true`) also stays eager so backward NaN checks run normally. A failure while recording raises instead of falling back to eager (RNG state and buffers are restored first), so disable `--model.capture_network` if recording fails on your model. Higher-order differentiation is not supported. Recording warmup restores RNG and BatchNorm buffers before the first real batch. Accumulation or custom optimiser/zero-grad hooks preserve live gradients outside the graph-owned backward buffers before they can be overwritten. AMP uses an uncached autocast scope as the PyTorch API requires. Validation, prediction and public `forward` calls stay eager. Checkpoints retain the ordinary parameter names and the flag, with no graph state. Each fit records afresh and drops its recording on completion or an exception. CUDA graph pools can increase reserved memory even when peak allocated memory falls; measure both on the intended workload. The measured scope and remaining validation gaps are in [WP-193c](work_packages/wp-193c.md). +For segmentation, the same `--model.capture_network true` option records only backbone and neck. Prediction groups, prototype fusion and both losses keep their existing arithmetic and gradient boundaries. It can combine with `--model.capture_segment_heads true`, which records the prediction groups separately. Gradient accumulation stays native, as do overridden or compiled segmentation forwards. The default remains `false`; a recording is transient and follows the same fit cleanup and input-signature rules above. See [WP-193g](work_packages/wp-193g.md) for the measured workload and limits. + ## Optional detection loss CUDA replay `--model.capture_loss true` replays the existing detection assignment and loss forward/backward with fresh predictions, targets and the current alpha tensor. Default `false`: capture costs setup time and retains CUDA workspace, so short or compute-bound runs may not benefit. It requires `task: detect`, single-device CUDA automatic optimisation and `compile_step: false`. It can combine with `capture_network`, `capture_orthogonalization` and `capture_updates`; `capture_loss` itself records no optimizer update, so the optimizer step is whatever those two optimizer options select (the existing eager path when both are off). diff --git a/docs/engineering/phase-17-pose-objective-parity.md b/docs/engineering/phase-17-pose-objective-parity.md index 3e8c37d..0186e85 100644 --- a/docs/engineering/phase-17-pose-objective-parity.md +++ b/docs/engineering/phase-17-pose-objective-parity.md @@ -26,6 +26,8 @@ --8<-- "work_packages/wp-193f.md:engineering" +--8<-- "work_packages/wp-193g.md:engineering" + --8<-- "work_packages/wp-194.md:engineering" --8<-- "work_packages/wp-195.md:engineering" diff --git a/docs/work_packages/wp-193d.md b/docs/work_packages/wp-193d.md index 2ee1eb9..302eae2 100644 --- a/docs/work_packages/wp-193d.md +++ b/docs/work_packages/wp-193d.md @@ -2,7 +2,7 @@ > Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) > -> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [WP-193f](../work_packages/wp-193f.md), [WP-200](../work_packages/wp-200.md) +> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [WP-193f](../work_packages/wp-193f.md), [WP-193g](../work_packages/wp-193g.md), [WP-200](../work_packages/wp-200.md) > > ⬅️ [WP-193c](../work_packages/wp-193c.md) · [Roadmap](../ROADMAP.md) · [WP-193e](../work_packages/wp-193e.md) ➡️ diff --git a/docs/work_packages/wp-193e.md b/docs/work_packages/wp-193e.md index 0135d21..4bf6dcd 100644 --- a/docs/work_packages/wp-193e.md +++ b/docs/work_packages/wp-193e.md @@ -2,7 +2,7 @@ > Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) > -> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [WP-193f](../work_packages/wp-193f.md) +> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [WP-193f](../work_packages/wp-193f.md), [WP-193g](../work_packages/wp-193g.md) > > ⬅️ [WP-193d](../work_packages/wp-193d.md) · [Roadmap](../ROADMAP.md) · [WP-193f](../work_packages/wp-193f.md) ➡️ diff --git a/docs/work_packages/wp-193f.md b/docs/work_packages/wp-193f.md index e06c857..442ae19 100644 --- a/docs/work_packages/wp-193f.md +++ b/docs/work_packages/wp-193f.md @@ -4,7 +4,7 @@ > > 🔗 Depends on [WP-193d](../work_packages/wp-193d.md), [WP-193e](../work_packages/wp-193e.md) · Needed by none > -> ⬅️ [WP-193e](../work_packages/wp-193e.md) · [Roadmap](../ROADMAP.md) · [WP-194](../work_packages/wp-194.md) ➡️ +> ⬅️ [WP-193e](../work_packages/wp-193e.md) · [Roadmap](../ROADMAP.md) · [WP-193g](../work_packages/wp-193g.md) ➡️ diff --git a/docs/work_packages/wp-193g.md b/docs/work_packages/wp-193g.md new file mode 100644 index 0000000..49e4b32 --- /dev/null +++ b/docs/work_packages/wp-193g.md @@ -0,0 +1,40 @@ +# 📦 WP-193g — `perf(ptl): replay shared segmentation features` + +> Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) +> +> 🔗 Depends on [WP-193d](../work_packages/wp-193d.md), [WP-193e](../work_packages/wp-193e.md) · Needed by none +> +> ⬅️ [WP-193f](../work_packages/wp-193f.md) · [Roadmap](../ROADMAP.md) · [WP-194](../work_packages/wp-194.md) ➡️ + + + +## WP-193g — replay common features without joining prediction gradients + + + +The segmentation prediction recorder leaves backbone and neck native. An L4 Nsight Systems trace of eight training batches, with prediction and optimizer replay already enabled, attributes 151.230 ms of host dispatch to backbone and 127.376 ms to neck within 1,544.688 ms of batch ranges. These diagnostic ranges account for 18.04% before their backward work; they are not an end-to-end speed claim. The optimizer-step range encloses the forward/backward closure, so its total cannot be interpreted as optimizer-only cost. + +The existing `capture_network` option now accepts segmentation and records these shared features through R43's existing graphed-callable adapter. All three feature levels already feed dense classification. Keeping coefficients, prototypes and semantic predictions outside that recording preserves their independent backward nodes and absent-gradient semantics. The option combines with `capture_segment_heads`, which records those predictions separately. Native model, assignment, loss, optimizer, schedule and EMA expressions are unchanged. + +Single-device CUDA automatic optimisation and a fixed input signature remain required. The default stays off. Accumulation, frozen shared features, altered head topology/dispatch or registered head hooks, overridden or compiled segmentation forwards and the shared recorder's existing unsupported conditions retain native execution. Partial batches do not create additional recordings. Transient recordings are released at fit completion or failure and do not alter checkpoint keys. + +Five alternating fresh-process pairs on an isolated NVIDIA L4 g2-standard-8 reduce complete nano segmentation fit time by 18.42% against main `b293970`, with segmentation prediction replay already enabled: + +| Seconds, median (range) | Native shared features | Replayed shared features | +| -- | -- | -- | +| Complete fit | 107.412 (106.764–109.816) | 87.629 (87.043–88.462) | +| Whole process | 112.721 (111.918–115.208) | 92.937 (92.471–93.827) | + +Each fit uses 512 public COCO val2017 images for training and 16 disjoint images for validation, the original 80-class annotations, five epochs, batch 4, 640 px and BF16 mixed precision. Both arms enable prepared-image caching, device HSV, epoch prefetch with four persistent workers, segmentation prediction replay, orthogonalization replay and update preparation replay. Mosaic remains active throughout (`close_mosaic=0`); warmup is one epoch. The timer starts before model, data and trainer construction and includes cache preparation, graph recording, training, EMA, validation and last-checkpoint writes. Imports and output hashing are outside that fit timer; the whole-process row includes them and reduces by 17.55%. Runtime: Python 3.11.14, torch 2.9.1+cu129, Lightning 2.6.6, driver 580.178.04. + +Model, optimizer, scheduler and EMA bytes, the 640 per-step loss bytes and validation metrics match within and across the ten fits. Losses, metrics and final model tensors are finite. Median whole-fit peak allocated memory is 1731 → 1511 MiB; peak reserved memory is 5496 → 5702 MiB. These are distinct allocator metrics, with extra reserved graph workspace rather than a capacity claim. The speed claim is limited to this batch-4 workload. + +The SHA-256 of the measured UTF-8 bytes of `src/lucid_yolo/ptl/module.py` at the measured PR-head tree (8ebf10e) is `1800875a7c034e6567ff8848e87b7c96cc091004fc112e1b5734c0f2a1354306`, computed with `sha256sum src/lucid_yolo/ptl/module.py`; later commits change `module.py`, so the hash does not describe later trees. The existing numerical stage expressions and all 41 restricted model/loss/assignment/optimizer/decode files are unchanged. Earlier measurements without the final head certificate are exploratory and do not supply this result. + +A separate single batch-16 screen (three epochs over the same 512 training and 16 validation images) changed complete fit time from 45.828 to 44.498 s, only 2.90%. Model, optimizer, scheduler, EMA, per-step losses and metrics matched. This is an exploratory boundary case, not a repeated batch-16 performance claim; no batch-16 optimization is submitted on that number. + +Focused checks pass 53 CPU-only cases with 43 GPU-marked cases deselected, 24 cases on torch 2.9.1 with CUDA available (including ten GPU-marked cases), and 27 cases including doctests on torch 2.4.0+cu124. Older BF16 replay retains native execution. The regressions exercise CPU routing, actual CUDA replay, every prediction and gradient byte, changing mask demand, empty supervision, partial batches, validation/EMA transitions, accumulation fallback and CLI construction. Removing replay deliberately fails the maintained route assertion. The full `make gate`, accelerator goldens and long-run accuracy were not run for this work and remain unverified; offline CI was green at the PR head 8ebf10e, which says nothing about the merged tree. + +No new method constant or numerical assumption is introduced. The basis is this repository's native stage composition, our CUDA profiles and R43's registered runtime API. + + diff --git a/docs/work_packages/wp-194.md b/docs/work_packages/wp-194.md index 4241df0..ec86030 100644 --- a/docs/work_packages/wp-194.md +++ b/docs/work_packages/wp-194.md @@ -4,7 +4,7 @@ > > 🔗 Depends on [WP-193](../work_packages/wp-193.md) · Needed by [WP-195](../work_packages/wp-195.md) > -> ⬅️ [WP-193f](../work_packages/wp-193f.md) · [Roadmap](../ROADMAP.md) · [WP-195](../work_packages/wp-195.md) ➡️ +> ⬅️ [WP-193g](../work_packages/wp-193g.md) · [Roadmap](../ROADMAP.md) · [WP-195](../work_packages/wp-195.md) ➡️ diff --git a/src/lucid_yolo/ptl/module.py b/src/lucid_yolo/ptl/module.py index eea1d30..4225066 100644 --- a/src/lucid_yolo/ptl/module.py +++ b/src/lucid_yolo/ptl/module.py @@ -227,6 +227,7 @@ _NetworkCaptureLifecycle, ) from lucid_yolo.ptl.seg_targets import instance_mask_targets, scale_boxes_to_grid, semantic_target +from lucid_yolo.ptl.segment_head_capture import _head_breach as _segment_head_breach from lucid_yolo.ptl.segment_head_capture import _SegmentHeadCapture, _SegmentHeadCaptureLifecycle if TYPE_CHECKING: @@ -778,11 +779,13 @@ def _validate_task_arguments(task: str, num_keypoints: int | None, rotated_iou_f ) -def _pose_head_topology(head: torch.nn.Module) -> tuple[tuple[object, ...], ...]: +def _head_topology(head: torch.nn.Module) -> tuple[tuple[object, ...], ...]: """Identify the head composition as constructed, whose dense outputs use every shared feature. The fingerprint is taken when the module is built, so the certificate it feeds holds the head unchanged since construction: a class patched before then is the reference. + Pose and segmentation both replay shared features into this head, so both compare + against the same fingerprint. """ return tuple( (name, type(module), type(module).forward, getattr(type(module), "forward_branch", None)) @@ -790,15 +793,36 @@ def _pose_head_topology(head: torch.nn.Module) -> tuple[tuple[object, ...], ...] ) +def _head_topology_breach(head: torch.nn.Module, topology: tuple[tuple[object, ...], ...], label: str) -> str | None: + """Name a head changed or hooked since construction, which replayed shared features cannot feed. + + Args: + head: The prediction head consuming the replayed shared features. + topology: The head's :func:`_head_topology` fingerprint taken at construction. + label: The task named in the reason, such as ``"pose"`` or ``"segmentation"``. + + Returns: + The retirement reason, or ``None`` while the head is still the one constructed. + """ + if _head_topology(head) != topology: + return f"the {label} head topology or class forward changed" + if _module_breach(head.modules()) is not None: + return f"a {label} head module is hooked, overridden or compiled" + return None + + def _pose_head_breach(head: torch.nn.Module, topology: tuple[tuple[object, ...], ...]) -> str | None: """Keep altered head control flow outside a recording of its shared inputs.""" if any(getattr(cls, name) is not native for cls, name, native in _POSE_STEM_DISPATCH): return "the native pose stem dispatch changed" - if _pose_head_topology(head) != topology: - return "the pose head topology or class forward changed" - if _module_breach(head.modules()) is not None: - return "a pose head module is hooked, overridden or compiled" - return None + return _head_topology_breach(head, topology, "pose") + + +#: Tasks whose network replay stays eager under gradient accumulation, with the label its +#: warning names. WP-193d reports the shared recorder's CUDA accumulation checks failing on +#: the unmodified baseline; detection is absent because whether it must stay eager too is +#: open in docs/ESCALATION.md (2026-10-20). +_ACCUMULATION_EAGER_LABELS: dict[str, str] = {"keypoints": "keypoint", "obb": "oriented", "segment": "segmentation"} def _validate_replay_options( @@ -806,21 +830,23 @@ def _validate_replay_options( ) -> None: """Reject replay options outside the objectives they record. - Network replay records the detection, pose or oriented stages; loss replay records the - detection objective only; segment-head replay records the segmentation prediction groups - only. Segment-head replay refuses the compiler although ``compile_step`` leaves the segment + Network replay records the detection, pose, oriented or segmentation stages; loss replay + records the detection objective only; segment-head replay records the segmentation prediction + groups only. Segment-head replay refuses the compiler although ``compile_step`` leaves the segment objective eager today, so an accepted configuration keeps its meaning once that objective is compiled. Raises: ValueError: If network replay is combined with a task other than detection, - keypoints or oriented detection, loss replay with a task other than detection, - segment-head replay with a task other than segmentation, or any of them with the - compiler (segment-head replay included, for forward compatibility with a compiled - segment objective). + keypoints, oriented detection or segmentation, loss replay with a task other than + detection, segment-head replay with a task other than segmentation, or any of them + with the compiler (segment-head replay included, for forward compatibility with a + compiled segment objective). """ - if capture_network and (task not in ("detect", "keypoints", "obb") or compile_step): - raise ValueError("capture_network=True requires task='detect', 'keypoints' or 'obb', and compile_step=False") + if capture_network and (task not in ("detect", "keypoints", "obb", "segment") or compile_step): + raise ValueError( + "capture_network=True requires task='detect', 'keypoints', 'obb' or 'segment', and compile_step=False" + ) if capture_loss and (task != "detect" or compile_step): raise ValueError("capture_loss=True requires task='detect' and compile_step=False") if capture_segment_heads and (task != "segment" or compile_step): @@ -1019,10 +1045,10 @@ class DetectionLitModule(LightningModule): 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. - capture_network: Replay fixed detection, keypoint or oriented 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"`` or ``task="obb"``, ``compile_step=False`` and + capture_network: 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 @@ -1035,9 +1061,13 @@ class DetectionLitModule(LightningModule): 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. Accumulation - or an overridden objective keeps it eager. Every task stays eager when the - LightningModule carries a hook, overrides - ``forward`` or is compiled. + or 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 also keeps accumulation native, + and 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 @@ -1065,8 +1095,8 @@ class DetectionLitModule(LightningModule): Raises: 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 or - oriented detection, ``capture_loss`` with a non-detection task, ``capture_segment_heads`` + if ``capture_network`` is combined with a task other than detection, keypoints, + oriented detection or segmentation, ``capture_loss`` with a non-detection task, ``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 @@ -1182,7 +1212,8 @@ def __init__( # noqa: PLR0913 — flat hyperparameter surface is deliberate: WP predict_angle=task == "obb", num_keypoints=num_keypoints if task == "keypoints" else None, ) - self._pose_head_topology = _pose_head_topology(self.head) if task == "keypoints" else () + #: The head as constructed, for the two tasks that replay shared features into it. + self._head_topology = _head_topology(self.head) if task in ("keypoints", "segment") else () self._init_segmentation(task, num_classes, capture_segment_heads) #: The keypoint objective, held as a submodule because R14's flow is the only #: loss in the project carrying parameters — hence a submodule rather than a @@ -1542,32 +1573,51 @@ def _detect_objective( ) return head_out, out - def _training_network(self, images: Tensor) -> DualHeadOutput: - """Run the dual head through the training recording when its contract holds. + def _network_replay_eligible(self, images: Tensor) -> bool: + """Decide whether a call may reach the network recording, before any task-specific contract. - Detection and oriented detection record all three stages; pose records shared - features and each stem independently. Unused point outputs have no backward - invocation, so batches with no assigned positives retain absent point-stem gradients. - All tasks keep changing instance counts and loss arithmetic outside capture. + Every task's training route shares this, so a change to when replay applies lands + once: only a training step of an activated fit, with automatic optimisation, enabled + gradients and CUDA images, qualifies. Everything else, evaluation included, stays native. + A retired recording answers every call eagerly and ignores later causes, so the + task-specific checks are skipped for the rest of that fit. """ - head_out: DualHeadOutput - if ( + return ( self._capture_network_active + and not self._network_capture.retired and self._trainer is not None and self._trainer.training and self.training and self.automatic_optimization and torch.is_grad_enabled() and images.is_cuda - ): - # WP-193d reports the shared recorder's CUDA accumulation checks failing on the unmodified - # baseline; whether detection must stay eager here too is open in docs/ESCALATION.md (2026-10-20). - if self._task in ("keypoints", "obb") and self.trainer.accumulate_grad_batches != 1: - label = "keypoint" if self._task == "keypoints" else "oriented" - self._network_capture.retire(f"{label} replay requires accumulate_grad_batches=1") + ) + + def _retire_accumulated_replay(self) -> None: + """Retire this task's network replay under gradient accumulation, which it was validated without. + + Tasks absent from :data:`_ACCUMULATION_EAGER_LABELS`, detection today, keep replaying. + ``stacklevel=4`` names the caller of the training route, as a retirement that route + decides itself does. + """ + label = _ACCUMULATION_EAGER_LABELS.get(self._task) + if label is not None and self.trainer.accumulate_grad_batches != 1: + self._network_capture.retire(f"{label} replay requires accumulate_grad_batches=1", stacklevel=4) + + def _training_network(self, images: Tensor) -> DualHeadOutput: + """Run the dual head through the training recording when its contract holds. + + Detection and oriented detection record all three stages; pose records shared + features and each stem independently. Unused point outputs have no backward + invocation, so batches with no assigned positives retain absent point-stem gradients. + All tasks keep changing instance counts and loss arithmetic outside capture. + """ + head_out: DualHeadOutput + if self._network_replay_eligible(images): + self._retire_accumulated_replay() if self._task == "keypoints": self._network_capture.check_stem_owners(self.head) - breach = _pose_head_breach(self.head, self._pose_head_topology) + breach = _pose_head_breach(self.head, self._head_topology) if breach is not None: self._network_capture.retire(breach) root_hooks = any( @@ -1774,7 +1824,7 @@ def forward_segmentation(self, images: Tensor) -> SegmentOutput: """ if self._task != "segment": raise ValueError(f"forward_segmentation requires task='segment'; this module's task is {self._task!r}") - features: tuple[Tensor, Tensor, Tensor] = self.neck(self.backbone(images)) + features = self._segmentation_features(images) runner = self._segment_head_capture.runner(self, features) detect = cast("DualHeadOutput", self.head(features) if runner is None else self.head(features, _runner=runner)) fused: Tensor = self.proto_fusion(features) @@ -1784,6 +1834,32 @@ def forward_segmentation(self, images: Tensor) -> SegmentOutput: detect=detect, prototypes=runner.mask(self.protonet, fused), semantic=runner.mask(self.semantic, fused) ) + def _segmentation_features(self, images: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Replay shared segmentation features while keeping mask supervision outside capture. + + The three feature levels always feed the dense classification objective. + Coefficients, prototypes and semantic predictions keep their independent + backward nodes, so an unused mask branch retains its absent gradients. + The existing holder owns precision/storage checks, buffer rollback and + the one-signature-per-fit lifecycle, as for shared pose features. + """ + if self._network_replay_eligible(images): + self._retire_accumulated_replay() + breach = _segment_head_breach(self.head) or _head_topology_breach( + self.head, self._head_topology, "segmentation" + ) + if breach is not None: + self._network_capture.retire(breach) + # Training calls `forward_segmentation` directly, never `self(...)`, so a compiled + # LightningModule call (`_compiled_call_impl`) never reaches this route; a compiled + # `forward_segmentation` assigned on the instance is caught by the `vars` check. + if type(self).forward_segmentation is not _EAGER_SEGMENT_FORWARD or "forward_segmentation" in vars(self): + self._network_capture.retire("the segmentation forward is overridden or compiled") + features = self._network_capture.run_features(images, (self.backbone, self.neck)) + if features is not None: + return features + return cast("tuple[Tensor, Tensor, Tensor]", self.neck(self.backbone(images))) + def forward_branch(self, images: Tensor, branch: BranchName) -> BranchOutput: """Run the backbone, neck, and **one** head branch over an image batch. @@ -2864,6 +2940,9 @@ def _anchor_grid(self, height: int, width: int, device: torch.device) -> tuple[T #: The unmodified public forward; instance/class overrides retain their eager dispatch. _EAGER_NETWORK_FORWARD = DetectionLitModule.forward +#: Segment composition overrides keep their native execution. +_EAGER_SEGMENT_FORWARD = DetectionLitModule.forward_segmentation + #: Consumers whose native calls retain gradients to every oriented output, even on empty batches. _ORIENTED_METHODS = tuple( (name, getattr(DetectionLitModule, name)) diff --git a/src/lucid_yolo/ptl/network_capture.py b/src/lucid_yolo/ptl/network_capture.py index 9a4c8f3..35a8ef1 100644 --- a/src/lucid_yolo/ptl/network_capture.py +++ b/src/lucid_yolo/ptl/network_capture.py @@ -3,8 +3,9 @@ The network has a fixed image signature while its ground-truth count changes per batch. Detection and oriented detection record backbone, neck and head; pose records -backbone, neck and each prediction stem separately. All three leave assignment, losses, -schedules, optimizer and EMA in their existing loop. PyTorch's graphed-callable API owns the +backbone, neck and each prediction stem separately; segmentation records backbone and neck. +All four leave assignment, losses, schedules, optimizer and EMA in their existing loop. +PyTorch's graphed-callable API owns the autograd bridge; no network arithmetic is rewritten. The recording is transient, bounded to one input signature and discarded at the fit boundary. @@ -123,11 +124,13 @@ def rebuild(tensors: tuple[Tensor, ...]) -> DualHeadOutput: class _FeatureNetwork(nn.Module): """Expose backbone and neck without fixing which optional head gradients exist. - A pose batch may have no assigned positives. Its point stems then keep - ``grad=None``; including those stems in one fixed backward would materialize - zero gradients and advance optimizer state that the native loop skips. - The shared features still feed the dense classification objective, while - the stems have independent autograd nodes and variable supervision stays native. + A pose or segmentation batch may have no assigned positives. Its optional + branches (pose point stems; segmentation coefficient, prototype and semantic + predictions) then keep ``grad=None``; including them in one fixed backward + would materialize zero gradients and advance optimizer state that the native + loop skips. The shared features still feed the dense classification objective, + while those branches keep independent backward nodes outside the recording and + variable supervision stays native. """ def __init__(self, backbone: nn.Module, neck: nn.Module) -> None: @@ -397,6 +400,29 @@ def __setstate__(self, state: dict[str, object]) -> None: self._label = None self.clear() + @property + def retired(self) -> bool: + """Report whether this fit's recording was retired, so every later call takes the eager route. + + Callers read it to skip their own replay checks once no cause they could find would + change the route; :meth:`clear` at the next fit resets it. + + Examples: + ```pycon + >>> import warnings + >>> holder = _NetworkCapture() + >>> holder.retired + False + >>> with warnings.catch_warnings(): + ... warnings.simplefilter("ignore") + ... holder.retire("a doctest cause") + >>> holder.retired + True + + ``` + """ + return self._disabled + def retire(self, reason: str, *, stacklevel: int = 3) -> None: """Retire a recording whose static execution contract no longer holds. @@ -452,9 +478,9 @@ def run( return self._run(images, stages, eager, lambda: _DetectionNetwork(*stages), _DetectionNetwork.rebuild) def run_features(self, images: Tensor, stages: tuple[nn.Module, nn.Module]) -> _Features | None: - """Replay shared pose features, or return no features so the full eager call stays observable.""" + """Replay shared features, or return no features so the full eager call stays observable.""" if not any(parameter.requires_grad for stage in stages for parameter in stage.parameters()): - self.retire("the shared pose features are frozen") + self.retire("the shared network features are frozen") return None def eager(_: Tensor) -> _Features | None: diff --git a/tests/ptl/test_network_capture.py b/tests/ptl/test_network_capture.py index fd99804..98e2ddf 100644 --- a/tests/ptl/test_network_capture.py +++ b/tests/ptl/test_network_capture.py @@ -87,22 +87,32 @@ def test_default_has_no_callback_or_registered_state() -> None: assert not candidate._capture_network_active +@pytest.mark.parametrize("task", ["detect", "keypoints", "obb", "segment"]) +def test_network_replay_accepts_all_supported_tasks_after_sync(task: str) -> None: + """Merging independent task extensions must retain every public constructor route.""" + options: dict[str, object] = {"num_keypoints": 5} if task == "keypoints" else {} + candidate = _module(task=task, capture_network=True, **options) + assert candidate.task == task + assert candidate.hparams.capture_network is True + assert not candidate._capture_network_active + + @pytest.mark.parametrize( "kwargs", [ - pytest.param({"task": "segment"}, id="segment"), pytest.param({"task": "obb", "compile_step": True}, id="obb-compile-step"), + pytest.param({"task": "segment", "compile_step": True}, id="segment-compile-step"), pytest.param({"compile_step": True}, id="detect-compile-step"), pytest.param({"task": "keypoints", "num_keypoints": 5, "compile_step": True}, id="keypoints-compile-step"), ], ) def test_incompatible_options_are_explicit(kwargs: dict[str, object]) -> None: - """Unsupported tasks and the separate compiler path fail before training begins. + """The separate compiler path fails before training begins. The whole message is matched, so it names the network-replay option and every task it accepts: a dropped or added task, or the loss-replay message, would fail here. """ - message = "capture_network=True requires task='detect', 'keypoints' or 'obb', and compile_step=False" + message = "capture_network=True requires task='detect', 'keypoints', 'obb' or 'segment', and compile_step=False" with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _module(capture_network=True, **kwargs) diff --git a/tests/ptl/test_pose_network_capture.py b/tests/ptl/test_pose_network_capture.py index 9271e03..271f954 100644 --- a/tests/ptl/test_pose_network_capture.py +++ b/tests/ptl/test_pose_network_capture.py @@ -457,23 +457,23 @@ def test_pose_cli_accepts_network_capture(capture: bool) -> None: def test_pose_head_certificate_survives_copy_and_rejects_overrides() -> None: """A native head stays eligible through deepcopy; changed dispatch cannot detach shared features.""" model = copy.deepcopy(_pose()) - assert _pose_head_breach(model.head, model._pose_head_topology) is None + assert _pose_head_breach(model.head, model._head_topology) is None with patch.object(model.head, "forward", wraps=model.head.forward): - assert _pose_head_breach(model.head, model._pose_head_topology) is not None - assert _pose_head_breach(model.head, model._pose_head_topology) is None + assert _pose_head_breach(model.head, model._head_topology) is not None + assert _pose_head_breach(model.head, model._head_topology) is None with patch.object(type(model.head), "forward", lambda self, values: values): - assert _pose_head_breach(model.head, model._pose_head_topology) is not None + assert _pose_head_breach(model.head, model._head_topology) is not None def test_pose_head_certificate_rejects_forward_branch_overrides() -> None: """``forward`` only reaches submodules through ``forward_branch``; both dispatch points count.""" model = copy.deepcopy(_pose()) - assert _pose_head_breach(model.head, model._pose_head_topology) is None + assert _pose_head_breach(model.head, model._head_topology) is None with patch.object(model.head, "forward_branch", wraps=model.head.forward_branch): - assert _pose_head_breach(model.head, model._pose_head_topology) is not None - assert _pose_head_breach(model.head, model._pose_head_topology) is None + assert _pose_head_breach(model.head, model._head_topology) is not None + assert _pose_head_breach(model.head, model._head_topology) is None with patch.object(type(model.head), "forward_branch", lambda self, features, branch: features): - assert _pose_head_breach(model.head, model._pose_head_topology) is not None + assert _pose_head_breach(model.head, model._head_topology) is not None @pytest.mark.parametrize( @@ -501,9 +501,9 @@ def test_altered_pose_head_dispatch_fails_certificate(mutate: Callable[[torch.nn Starts from a fresh deepcopy certified clean, applies one mutation, and expects a named breach. """ model = copy.deepcopy(_pose()) - assert _pose_head_breach(model.head, model._pose_head_topology) is None + assert _pose_head_breach(model.head, model._head_topology) is None mutate(model.head) - assert _pose_head_breach(model.head, model._pose_head_topology) is not None + assert _pose_head_breach(model.head, model._head_topology) is not None def test_frozen_features_leave_full_native_call_to_owner() -> None: @@ -512,7 +512,7 @@ def test_frozen_features_leave_full_native_call_to_owner() -> None: model.backbone.requires_grad_(False) model.neck.requires_grad_(False) holder = _NetworkCapture() - with pytest.warns(UserWarning, match="shared pose features are frozen"): + with pytest.warns(UserWarning, match="shared network features are frozen"): features = holder.run_features(torch.zeros(2, 3, 64, 64), (model.backbone, model.neck)) assert features is None and holder._disabled and holder._network is None @@ -650,6 +650,30 @@ def test_pose_routing_retires_then_runs_native( assert output.o2m_cls.shape[0] == 2 +def test_retired_holder_skips_pose_replay_predicates() -> None: + """A retired recording sends later steps through the full native call without re-checking the head. + + Once retired, the owner check and every later ``retire`` are no-ops, so the stem-owner + walk and the head certificate could change nothing; a step after retirement skips both + and never reaches the recorder. + """ + torch.manual_seed(42) + model = _routed_pose() + with pytest.warns(UserWarning, match="stays eager: a test cause"): + model._network_capture.retire("a test cause") + images = _cuda_flagged_images() + with ( + patch("lucid_yolo.ptl.module._pose_head_breach", autospec=True) as breach, + patch.object(model._network_capture, "check_stem_owners") as owners, + patch.object(model._network_capture, "run_features") as run, + ): + output = model._training_network(images) + breach.assert_not_called() + owners.assert_not_called() + run.assert_not_called() + assert output.o2m_cls.shape[0] == 2 + + def test_pose_routing_falls_back_when_features_are_absent() -> None: """When the recorder declines to return features, the full native call runs without a warning.""" torch.manual_seed(42) @@ -818,7 +842,7 @@ def custom(self: torch.nn.Module, *args: object, **kwargs: object) -> object: with patch.object(owner, name, custom): model = _pose() - assert _pose_head_breach(model.head, model._pose_head_topology) == "the native pose stem dispatch changed" + assert _pose_head_breach(model.head, model._head_topology) == "the native pose stem dispatch changed" def test_stem_holder_dispatch_and_lifetime_without_cuda() -> None: diff --git a/tests/ptl/test_segment_network_capture.py b/tests/ptl/test_segment_network_capture.py new file mode 100644 index 0000000..1422016 --- /dev/null +++ b/tests/ptl/test_segment_network_capture.py @@ -0,0 +1,641 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared segmentation replay keeps prediction gradients and the fit lifecycle native.""" + +from __future__ import annotations + +import copy +import warnings +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from dataclasses import fields +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from pytorch_lightning import Callback, Trainer +from pytorch_lightning.callbacks import GradientAccumulationScheduler +from torch.utils.data import DataLoader + +from lucid_yolo.cli.train import DetectionCLI, packaged_config +from lucid_yolo.data.targets import Targets +from lucid_yolo.models.build import SegmentOutput +from lucid_yolo.models.heads.detect import DualHeadOutput +from lucid_yolo.ptl.callbacks import EMACallback +from lucid_yolo.ptl.datamodule import DetectionDataModule +from lucid_yolo.ptl.module import DetectionLitModule +from lucid_yolo.ptl.network_capture import _BF16_REPLAY_SUPPORTED + + +@pytest.fixture(autouse=True) +def _deterministic(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + """Restore the caller's deterministic policy and cudnn benchmark flag after each test. + + Lightning's trainer setup writes `torch.backends.cudnn.benchmark` unconditionally, so a fit test would leak it. + """ + monkeypatch.setattr(torch.backends.cudnn, "benchmark", torch.backends.cudnn.benchmark) + enabled = torch.are_deterministic_algorithms_enabled() + warning = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warning) + + +def _module(**kwargs: object) -> DetectionLitModule: + """Build the real segmentation stages at a small test width. + + Examples: + >>> _module().task + 'segment' + """ + return DetectionLitModule(depth=0.34, width=0.25, max_channels=64, num_classes=2, task="segment", **kwargs) + + +@contextmanager +def _no_retire_warning() -> Iterator[None]: + """Fail on either replay holder reporting that it stays eager. + + Examples: + >>> with _no_retire_warning(): + ... pass + """ + with warnings.catch_warnings(): + warnings.filterwarnings("error", message=".*stays eager", category=UserWarning) + yield + + +def _same(left: object, right: object) -> None: + """Compare tensor bytes and the complete nested optimizer state. + + Examples: + >>> _same(torch.tensor([1.0]), torch.tensor([1.0])) + """ + if isinstance(left, torch.Tensor): + assert isinstance(right, torch.Tensor) + assert left.dtype == right.dtype and left.shape == right.shape + assert torch.equal( + left.detach().contiguous().reshape(-1).view(torch.uint8), + right.detach().contiguous().reshape(-1).view(torch.uint8), + ) + elif isinstance(left, dict): + assert isinstance(right, dict) and left.keys() == right.keys() + for key in left: + _same(left[key], right[key]) + elif isinstance(left, (tuple, list)): + assert isinstance(right, (tuple, list)) and len(left) == len(right) + for first, second in zip(left, right, strict=True): + _same(first, second) + else: + assert left == right + + +def _batch(empty: bool = False, size: int = 2) -> tuple[torch.Tensor, list[Targets], list[torch.Tensor]]: + """Pair real image tensors with explicit instance masks on the prototype grid. + + Examples: + >>> images, targets, masks = _batch(True, 1) + >>> images.shape, targets[0].boxes.shape, masks[0].shape + (torch.Size([1, 3, 64, 64]), torch.Size([0, 4]), torch.Size([0, 16, 16])) + """ + images = torch.rand(size, 3, 64, 64) + ring = torch.tensor([[8.0, 8.0], [48.0, 8.0], [48.0, 48.0], [8.0, 48.0]]) + targets = [ + Targets(boxes=torch.empty(0, 4), labels=torch.empty(0, dtype=torch.int64)) + if empty + else Targets(boxes=torch.tensor([[8.0, 8.0, 48.0, 48.0]]), labels=torch.tensor([1]), polygons=[ring]) + for _ in range(size) + ] + masks = [torch.zeros(0 if empty else 1, 16, 16) for _ in range(size)] + for mask in masks: + mask[:, 2:12, 2:12] = 1.0 + return images, targets, masks + + +def test_segment_request_preserves_defaults_and_checkpoint_keys() -> None: + """A requested feature recorder adds no parameters and stays inactive off CUDA.""" + native, replay = _module(), _module(capture_network=True) + assert not native.hparams.capture_network + assert native.state_dict().keys() == replay.state_dict().keys() + assert len(replay.configure_callbacks()) == 1 + with pytest.warns(UserWarning, match="single-device CUDA"): + replay.setup("fit") + with patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("CPU must stay native")): + assert replay.forward_segmentation(torch.rand(2, 3, 64, 64)).prototypes.shape == (2, 32, 16, 16) + assert replay._network_capture._network is None + + +@pytest.mark.parametrize( + ("heads", "dtype"), + [ + pytest.param(False, torch.float32, id="network-fp32"), + pytest.param(False, torch.float16, id="network-fp16"), + pytest.param(False, torch.bfloat16, id="network-bf16"), + pytest.param(True, torch.bfloat16, id="heads-bf16"), + ], +) +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA graph execution") +def test_segment_shared_replay_preserves_predictions_and_unused_gradients(heads: bool, dtype: torch.dtype) -> None: + """Changing mask demand keeps every prediction, gradient and BatchNorm byte native.""" + torch.manual_seed(347) + native = _module(capture_segment_heads=heads).cuda().to(memory_format=torch.channels_last) + replay = copy.deepcopy(native) + replay._capture_network_active = True + replay._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + native._segment_head_capture.active = replay._segment_head_capture.active = heads + native._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + for use_masks in (True, False, True): + images = torch.rand(2, 3, 64, 64, device="cuda").to(memory_format=torch.channels_last) + outputs = [] + for model in (native, replay): + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32): + result = model.forward_segmentation(images) + tensors = [result.detect.o2m_cls, result.detect.o2o_cls] + if use_masks: + tensors += [result.prototypes, result.semantic, result.detect.o2m_coeff, result.detect.o2o_coeff] + loss = sum(t.float().square().mean() for t in tensors if t is not None) + loss.backward() + outputs.append(result) + for field in fields(DualHeadOutput): + _same(getattr(outputs[0].detect, field.name), getattr(outputs[1].detect, field.name)) + _same(outputs[0].prototypes, outputs[1].prototypes) + _same(outputs[0].semantic, outputs[1].semantic) + _same(native.state_dict(), replay.state_dict()) + unused = 0 + for (name, left), (_, right) in zip(native.named_parameters(), replay.named_parameters(), strict=True): + _same(left.grad, right.grad) + if not use_masks and ( + "coeff_stems" in name or name.startswith(("protonet.", "semantic.", "proto_fusion.")) + ): + assert left.grad is None and right.grad is None + unused += 1 + assert use_masks or unused > 0, "no mask-only parameter matched the unused-gradient filter" + assert (replay._network_capture._network is not None) == (dtype != torch.bfloat16 or _BF16_REPLAY_SUPPORTED) + + +#: Retirement warning of the network holder under gradient accumulation, which segmentation replay lacks validation for. +_ACCUMULATION_RETIRE = "capture_network=True stays eager: segmentation replay requires accumulate_grad_batches=1" +#: Retirement warning of the head holder outside BF16 mixed precision. +_HEAD_PRECISION_RETIRE = "capture_segment_heads=True stays eager" + + +@pytest.mark.parametrize( + ("heads", "precision", "accumulation", "retire"), + [ + pytest.param(False, "bf16-mixed", 1, None, id="network-bf16"), + pytest.param(False, "16-mixed", 1, None, id="network-fp16"), + pytest.param(False, "32-true", 1, None, id="network-fp32"), + pytest.param(True, "bf16-mixed", 1, None, id="heads-bf16"), + pytest.param(True, "bf16-mixed", 2, _ACCUMULATION_RETIRE, id="heads-bf16-accumulated"), + pytest.param(True, "16-mixed", 1, _HEAD_PRECISION_RETIRE, id="heads-fp16-stay-eager"), + ], +) +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires public CUDA training") +def test_segment_fit_replays_across_empty_partial_validation_and_ema( + heads: bool, precision: str, accumulation: int, retire: str | None +) -> None: + """Real fits preserve MuSGD and EMA through validation and unsupported accumulation or precision. + + Network-only cells cover every precision; head replay adds the supported bf16 cell, the accumulation + retirement and the fp16 retirement, whose warning must name the holder that stays eager. + """ + torch.manual_seed(23) + batches = [_batch(), _batch(True), _batch(size=1)] + initial = _module().state_dict() + snapshots = [] + recordings = [] + + class Observe(Callback): + """Observe the live feature recording before fit cleanup releases it.""" + + def on_train_batch_end( + self, trainer: Trainer, pl_module: DetectionLitModule, outputs: object, batch: object, batch_idx: int + ) -> None: + """Record whether the public route owns a real graphed network.""" + recordings.append(pl_module._network_capture._network is not None) + + for capture in (False, True): + torch.manual_seed(23) + model = _module(capture_network=capture, capture_segment_heads=heads, lr=1e-4, warmup_epochs=0.0) + model.load_state_dict(initial) + ema = EMACallback() + trainer = Trainer( + accelerator="cuda", + devices=1, + precision=precision, + max_epochs=2, + deterministic=True, + benchmark=False, + accumulate_grad_batches=accumulation, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + num_sanity_val_steps=0, + callbacks=[ema, Observe()], + gradient_clip_val=1.0, + ) + replays = precision != "bf16-mixed" or _BF16_REPLAY_SUPPORTED + combined = capture and heads and accumulation == 1 and precision == "bf16-mixed" and _BF16_REPLAY_SUPPORTED + expectation = ( + pytest.warns(UserWarning, match=retire) + if capture and retire + else _no_retire_warning() + if capture and replays + else nullcontext() + ) + with ( + expectation, + patch("torch.cuda.make_graphed_callables", wraps=torch.cuda.make_graphed_callables) as record, + ): + trainer.fit( + model, + train_dataloaders=DataLoader(batches, batch_size=None), + val_dataloaders=DataLoader([_batch()], batch_size=None), + ) + assert trainer.state.finished + feature_records = [call for call in record.call_args_list if type(call.args[0]).__name__ == "_FeatureNetwork"] + assert len(feature_records) == int(capture and accumulation == 1 and replays) + head_records = [call for call in record.call_args_list if type(call.args[0]).__name__ == "_PredictionGroup"] + assert len(head_records) >= int(combined) + scaler = getattr(trainer.precision_plugin, "scaler", None) + snapshots.append( + ( + copy.deepcopy(model.state_dict()), + copy.deepcopy(trainer.optimizers[0].state_dict()), + copy.deepcopy(ema.state_dict()), + copy.deepcopy(scaler.state_dict()) if scaler is not None else None, + ) + ) + assert model._network_capture._network is None + assert model._segment_head_capture.entries == {} + _same(*snapshots) + assert recordings[:6] == [False] * 6 + assert recordings[6:] == [accumulation == 1 and replays] * 6 + + +class _CudaMarked(torch.Tensor): + """Exercise host routing with real CPU storage and a CUDA eligibility flag.""" + + @property + def is_cuda(self) -> bool: + """Mark only the dispatch boundary; tensor arithmetic remains on CPU.""" + return True + + +#: Modes whose call stays native before the recorder is reached, so `run_features` is never consulted. +_NATIVE_BEFORE_RECORDER = ( + "inactive", + "no-trainer", + "validating", + "manual-optimization", + "cpu-images", + "eval", + "no-grad", +) + + +#: Reason completing the `capture_network=True stays eager: ` warning of each retiring mode; others run warning-free. +_RETIRE_REASON = { + "accumulate": "segmentation replay requires accumulate_grad_batches=1$", + "override": "the segmentation forward is overridden or compiled$", + "head-hook": "a segmentation head module is hooked, overridden or compiled$", + "head-helper": "a prediction branch overrides its native dispatch$", +} + + +@pytest.mark.parametrize( + "mode", + [ + "record", + "absent", + "accumulate", + "override", + "compiled", + "head-hook", + "head-helper", + "eval", + "no-grad", + "inactive", + "no-trainer", + "validating", + "manual-optimization", + "cpu-images", + ], +) +def test_segment_host_routing_and_native_fallback(mode: str) -> None: + """A real segmentation forward consumes replayed features or invokes both native stages.""" + model = _module(capture_network=True) + model._capture_network_active = mode != "inactive" + model.automatic_optimization = mode != "manual-optimization" + model._trainer = ( + None + if mode == "no-trainer" + else SimpleNamespace(training=mode != "validating", accumulate_grad_batches=2 if mode == "accumulate" else 1) + ) + images = torch.rand(2, 3, 64, 64) + images = images if mode == "cpu-images" else images.as_subclass(_CudaMarked) + features = model.neck(model.backbone(images)) + if mode == "override": + model.forward_segmentation = model.forward_segmentation + elif mode == "compiled": + model._compiled_call_impl = lambda *_: None + elif mode == "head-hook": + model.head.register_forward_hook(lambda *_: None) + elif mode == "head-helper": + model.head.o2m._base_predictions = model.head.o2m._base_predictions + elif mode == "eval": + model.eval() + expected_call = mode not in _NATIVE_BEFORE_RECORDER + context = torch.no_grad() if mode == "no-grad" else nullcontext() + reason = _RETIRE_REASON.get(mode) + warning = ( + pytest.warns(UserWarning, match=f"capture_network=True stays eager: {reason}") + if reason + else warnings.catch_warnings(action="error") + ) + with ( + context, + warning, + patch.object( + model._network_capture, "run_features", return_value=features if mode == "record" else None + ) as run, + patch.object(model.backbone, "forward", wraps=model.backbone.forward) as backbone, + ): + actual = model.forward_segmentation(images) + assert run.call_count == int(expected_call) + assert backbone.call_count == int(mode != "record") + assert model._network_capture.retired == (reason is not None) + assert actual.prototypes.shape == (2, 32, 16, 16) + if mode == "record": + expected = model.head(features) + for field in fields(DualHeadOutput): + _same(getattr(actual.detect, field.name), getattr(expected, field.name)) + + +class _OverridingModule(DetectionLitModule): + """A subclass whose class-level `forward_segmentation` replaces the native one.""" + + def forward_segmentation(self, images: torch.Tensor) -> SegmentOutput: + """Delegate to the native forward, which the recorder can no longer certify by identity.""" + return super().forward_segmentation(images) + + +def _eligible_module(cls: type[DetectionLitModule] = DetectionLitModule) -> DetectionLitModule: + """Build an activated segmentation fit whose images are eligible for replay. + + Examples: + >>> _eligible_module()._capture_network_active + True + """ + model = cls(depth=0.34, width=0.25, max_channels=64, num_classes=2, task="segment", capture_network=True) + model._capture_network_active = True + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + return model + + +def test_segment_class_level_forward_override_retires_to_native() -> None: + """A subclass overriding `forward_segmentation` on the class retires replay and stays native.""" + model = _eligible_module(_OverridingModule) + images = torch.rand(2, 3, 64, 64).as_subclass(_CudaMarked) + with ( + pytest.warns(UserWarning, match="stays eager: the segmentation forward is overridden or compiled$"), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("overridden forward must stay native")), + ): + actual = model.forward_segmentation(images) + assert model._network_capture.retired + assert actual.prototypes.shape == (2, 32, 16, 16) + + +def test_segment_head_topology_swap_retires_with_topology_reason() -> None: + """Swapping the class of one head submodule retires replay naming the changed topology.""" + model = _eligible_module() + stem = model.head.o2m.box_stems[0] + stem.__class__ = type("_SwappedStem", (type(stem),), {}) + images = torch.rand(2, 3, 64, 64).as_subclass(_CudaMarked) + with ( + pytest.warns(UserWarning, match="stays eager: the segmentation head topology or class forward changed$"), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("swapped head must stay native")), + ): + actual = model.forward_segmentation(images) + assert model._network_capture.retired + assert actual.prototypes.shape == (2, 32, 16, 16) + + +def test_retired_holder_skips_segment_replay_predicates() -> None: + """A retired recording sends later steps native without re-walking the head certificate. + + Once retired, the holder answers every call eagerly and ignores later causes, so the + dispatch and topology walks could change nothing; a step after retirement skips them + and never reaches the recorder. + """ + model = _module(capture_network=True) + model._capture_network_active = True + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + with pytest.warns(UserWarning, match="stays eager: a test cause"): + model._network_capture.retire("a test cause") + images = torch.rand(2, 3, 64, 64).as_subclass(_CudaMarked) + with ( + warnings.catch_warnings(), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("retired holder must stay native")), + patch("lucid_yolo.ptl.module._segment_head_breach", autospec=True) as dispatch, + patch("lucid_yolo.ptl.module._head_topology_breach", autospec=True) as topology, + patch.object(model._network_capture, "run_features") as run, + ): + warnings.simplefilter("error") + features = model._segmentation_features(images) + dispatch.assert_not_called() + topology.assert_not_called() + run.assert_not_called() + assert len(features) == 3 + assert model._network_capture.retired + + +def test_frozen_segment_features_retire_before_cuda_submission() -> None: + """Frozen features stay native after unfreezing and report the shared-stage cause.""" + model = _module(capture_network=True) + model._capture_network_active = True + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + model.backbone.requires_grad_(False) + model.neck.requires_grad_(False) + images = torch.rand(2, 3, 64, 64).as_subclass(_CudaMarked) + with ( + pytest.warns(UserWarning, match="shared network features are frozen"), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("frozen stages must stay native")), + ): + model.forward_segmentation(images) + assert model._network_capture._disabled + model.backbone.requires_grad_(True) + model.neck.requires_grad_(True) + with ( + warnings.catch_warnings(), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("retirement must stay sticky")), + ): + warnings.simplefilter("error") + model.forward_segmentation(images) + assert model._network_capture._disabled + assert model._network_capture._network is None + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_segmentation_cli_keeps_feature_and_prediction_options_independent(enabled: bool) -> None: + """The shipped segmentation recipe forwards both requests through the real CLI.""" + cli = DetectionCLI( + DetectionLitModule, + DetectionDataModule, + args=[ + "--config", + str(packaged_config("seg_nano_smoke.yaml")), + "--model.capture_network", + str(enabled).lower(), + "--model.capture_segment_heads", + "true", + "--trainer.accelerator", + "cpu", + "--trainer.logger", + "false", + ], + run=False, + save_config_callback=None, + ) + assert cli.model.task == "segment" + assert cli.model.hparams.capture_network is enabled + assert cli.model.hparams.capture_segment_heads is True + assert cli.datamodule._mask_targets + + +def test_segment_head_hook_retires_feature_replay_and_stays_observable() -> None: + """A mutable head cannot change which shared outputs a recorded backward must consume.""" + model = _module(capture_network=True) + model._capture_network_active = True + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=1) + images = torch.rand(2, 3, 64, 64).as_subclass(_CudaMarked) + observed = [] + handle = model.head.register_forward_hook(lambda *_: observed.append(True)) + with ( + pytest.warns(UserWarning, match="stays eager: a segmentation head module is hooked, overridden or compiled$"), + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("custom head must stay native")), + ): + model.forward_segmentation(images) + assert observed == [True] and model._network_capture._disabled + handle.remove() + model.forward_segmentation(images) + assert model._network_capture._network is None + + +#: Trainer arguments shared by the CUDA lifecycle fits, which differ only in precision and callbacks. +_LIFECYCLE_FIT = { + "accelerator": "cuda", + "devices": 1, + "max_epochs": 2, + "deterministic": True, + "benchmark": False, + "logger": False, + "enable_checkpointing": False, + "enable_progress_bar": False, + "enable_model_summary": False, + "num_sanity_val_steps": 0, +} +_skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires public CUDA training") +_skip_no_bf16_replay = pytest.mark.skipif(not _BF16_REPLAY_SUPPORTED, reason="BF16 replay requires torch>=2.9") + + +class _FailingModule(DetectionLitModule): + """A segmentation module whose second training step fails after the first recorded its graphs.""" + + def training_step(self, batch: object, batch_idx: int) -> torch.Tensor: + """Run the native step, then fail on the second batch of the fit.""" + if batch_idx == 1: + raise RuntimeError("step failed after recording") + return super().training_step(batch, batch_idx) # type: ignore[arg-type] + + +class _FreezeBackbone(Callback): + """Freeze the shared backbone after the first batch and release it after the second.""" + + def on_train_batch_end( + self, trainer: Trainer, pl_module: DetectionLitModule, outputs: object, batch: object, batch_idx: int + ) -> None: + """Change the trainable set of a recorded network between steps.""" + pl_module.backbone.requires_grad_(batch_idx != 0) + + +def _fit_batches() -> DataLoader[object]: + """Serve full, empty-mask and partial batches, as the fit parity test does. + + Examples: + >>> len(_fit_batches()) + 3 + """ + return DataLoader([_batch(), _batch(True), _batch(size=1)], batch_size=None) # type: ignore[arg-type] + + +@pytest.mark.gpu +@_skip_no_cuda +@_skip_no_bf16_replay +def test_segment_fit_exception_after_recording_releases_both_holders() -> None: + """A failing step after recording leaves no graph, pending backward or head entry behind.""" + torch.manual_seed(23) + model = _FailingModule( + depth=0.34, + width=0.25, + max_channels=64, + num_classes=2, + task="segment", + capture_network=True, + capture_segment_heads=True, + ) + trainer = Trainer(**_LIFECYCLE_FIT, precision="bf16-mixed") + with ( + patch("torch.cuda.make_graphed_callables", wraps=torch.cuda.make_graphed_callables) as record, + pytest.raises(RuntimeError, match="step failed after recording"), + ): + trainer.fit(model, train_dataloaders=_fit_batches()) + assert any(type(call.args[0]).__name__ == "_FeatureNetwork" for call in record.call_args_list) + assert model._network_capture._network is None + assert model._network_capture._pending_backward is False + assert model._segment_head_capture.entries == {} + + +@pytest.mark.gpu +@_skip_no_cuda +def test_segment_second_fit_records_once_per_fit() -> None: + """Fitting the same module twice records its shared features once in each fit.""" + torch.manual_seed(23) + model = _module(capture_network=True) + counts = [] + for _ in range(2): + trainer = Trainer(**_LIFECYCLE_FIT, precision="16-mixed") + with patch("torch.cuda.make_graphed_callables", wraps=torch.cuda.make_graphed_callables) as record: + trainer.fit(model, train_dataloaders=_fit_batches()) + counts.append(sum(type(call.args[0]).__name__ == "_FeatureNetwork" for call in record.call_args_list)) + assert counts == [1, 1] + + +@pytest.mark.gpu +@_skip_no_cuda +def test_segment_freeze_after_recording_retires_once() -> None: + """Freezing then releasing the backbone after recording reports a single retirement.""" + torch.manual_seed(23) + model = _module(capture_network=True) + trainer = Trainer(**_LIFECYCLE_FIT, precision="16-mixed", callbacks=[_FreezeBackbone()]) + with pytest.warns(UserWarning, match="capture_network=True stays eager") as record: + trainer.fit(model, train_dataloaders=_fit_batches()) + assert sum("capture_network=True stays eager" in str(warning.message) for warning in record) == 1 + + +@pytest.mark.gpu +@_skip_no_cuda +def test_segment_accumulation_change_after_recording_retires_through_holder() -> None: + """Raising accumulate_grad_batches after epoch one retires the recording with the accumulation reason.""" + torch.manual_seed(23) + model = _module(capture_network=True) + scheduler = GradientAccumulationScheduler({0: 1, 1: 2}) + trainer = Trainer(**_LIFECYCLE_FIT, precision="16-mixed", callbacks=[scheduler]) + with pytest.warns(UserWarning, match=_ACCUMULATION_RETIRE): + trainer.fit(model, train_dataloaders=_fit_batches()) + assert model._network_capture._network is None