diff --git a/CHANGELOG.md b/CHANGELOG.md index ef55422..8ea9e46 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 +- **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. - **Opt-in shared-feature CUDA replay for pose training.** `--model.capture_network true` can record the existing backbone and neck for `task: keypoints`; the head stays native so a batch with no assigned positives keeps absent point-stem gradients. Accumulated pose runs, runs with a fully frozen backbone and neck, and heads with changed topology, hooks or forwards stay native. The default remains off. The speed claim is limited to a small-batch workload (batch 4, one L4): a single batch-16 pair changed fit time by 2.50%, and replay lowered peak allocated memory while raising peak reserved memory, so check throughput and both memory figures on your own batch size. See [WP-193d](docs/work_packages/wp-193d.md) for complete-fit measurements and validation limits. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index b04c908..a66ddd1 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -1,6 +1,6 @@ # 🗺️ Work-Package Roadmap -The agent's work queue: 199 numbered work packages, one commit each, executed in dependency order per the AGENTS.md loop, plus lettered rows recording follow-up work that landed after its package had closed. Transcribed from the governing blueprint (sec. 15) with a live status column. +The agent's work queue: 200 numbered work packages, one commit each, executed in dependency order per the AGENTS.md loop, plus lettered rows recording follow-up work that landed after its package had closed. Transcribed from the governing blueprint (sec. 15) with a live status column. **Scope says what a package does, not what executing it taught.** Measurements, rejected approaches and negative results live in RESEARCH_LOG.md, linked per row as `· log`; choices the papers left open live in ASSUMPTIONS.md and are cited by id; release results live in REPRODUCTION_REPORT.md. A row that restates any of those is duplicating a record that has an owner. @@ -364,6 +364,7 @@ Opened by issue 2 on the tracker: R1 sec. 3.4.2 publishes the pose objective thi | 197c | `perf(optim): replay CUDA update preparation` [DATA][GPU] | Opt-in group preparation replay retains native arithmetic, records momentum with rollback-safe setup and applies changing learning rates outside capture; layout alignment enables final foreach writes (D29) · [log](work_packages/wp-197c.md#wp-197c) | `tests/optim/test_update_replay.py`; `tests/ptl/test_cli.py::test_capture_updates_flag_reaches_optimizer`; focused optimizer/replay/CLI CPU/CUDA checks, mutation proof, five paired public-COCO fits with byte parity and 10.01% median complete-fit reduction (nano, batch 4, BF16, L4; batch-32 screen showed no benefit; batch 16 unmeasured); mypy and full precommit; full `make gate` not run | 197b | 🔄 | | 198 | `perf(data): overlap epoch preparation and fuse colour restore` [DATA][GPU] | Persistent workers prepare up to `num_workers x prefetch_factor` next-epoch batches during validation; an optional Triton pass restores the existing colour boundary without intermediate batches. Fixed sampler ownership, finite epochs and exact pixel codes remain the contract · [log](work_packages/wp-198.md#wp-198) | `tests/ptl/test_epoch_prefetch.py`; `tests/ptl/test_fused_device_hsv.py`; isolated-L4 complete-fit comparisons; full `make gate` not run | 197 | 🔄 | | 199 | `perf(ptl): replay detection assignment and loss` [DATA][GPU] | Opt-in loss forward/backward replay calls the existing assignment and scoring expressions with masked target capacity. One recording grows from observed counts; checks and live alpha stay current, and optimizer work is independent · [log](work_packages/wp-199.md#wp-199) | `tests/ptl/test_loss_capture.py`; focused CPU/CUDA, mutation checks and repeated complete public-COCO fits; full `make gate` and accelerator goldens not run | 198 | 🔄 | +| 200 | `perf(ptl): replay native CUDA oriented networks` [DATA][GPU] | Opt-in network replay retains all six native oriented outputs, while classification, rotated objectives and state updates stay native. Accumulation and changed objectives retain the original route · [log](work_packages/wp-200.md#wp-200) | `tests/ptl/test_oriented_network_capture.py`; actual graph execution, empty-angle gradients and validation transitions; five paired public-DOTA complete fits above 8%; focused CPU/CUDA and full pre-commit; full `make gate`, accelerator goldens and long-run accuracy not run | 193d, 199 | 🔄 | ## ⛓️ Critical path and parallelism diff --git a/docs/TRAINING.md b/docs/TRAINING.md index 2f1d4aa..4e341e7 100644 --- a/docs/TRAINING.md +++ b/docs/TRAINING.md @@ -143,7 +143,9 @@ Budget about `12 * sum(height * width)` bytes for the retained RGB pixels, plus For `task: keypoints`, the same flag records only backbone and neck; the complete dual head and the current RLE/OKS or Laplace objectives stay native. That preserves absent point-stem gradients on batches with no assigned positives. Gradient accumulation, a fully frozen backbone and neck, changed head topology or forwards, and head hooks/compilation keep pose native; a partial freeze that leaves any backbone or neck parameter trainable still replays. The trainable set is fixed per fit at the first recording, so a fully frozen start, or a later freeze or unfreeze of any backbone or neck parameter, keeps the rest of the fit native after one warning: a recipe that starts with backbone and neck fully frozen trains natively throughout, even after it unfreezes them. The head's ordinary dense classification path supplies gradients to all three shared features. The pose speed-up is likewise limited to a small-batch workload (batch 4): a single batch-16 pair changed fit time by 2.50%, and replay lowered peak allocated memory while raising peak reserved memory (see the memory note below), so check throughput and memory on your own batch size. The pose measurements and remaining validation limits are recorded in [WP-193d](work_packages/wp-193d.md). -This path requires single-device CUDA automatic optimisation and `task: detect` or `task: keypoints`. 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 are checked on 2.4 as well. 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. +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 are checked on 2.4 as well. 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). diff --git a/docs/engineering/phase-17-pose-objective-parity.md b/docs/engineering/phase-17-pose-objective-parity.md index 9c52cc9..13ce4b4 100644 --- a/docs/engineering/phase-17-pose-objective-parity.md +++ b/docs/engineering/phase-17-pose-objective-parity.md @@ -37,3 +37,5 @@ --8<-- "work_packages/wp-198.md:engineering" --8<-- "work_packages/wp-199.md:engineering" + +--8<-- "work_packages/wp-200.md:engineering" diff --git a/docs/work_packages/wp-193d.md b/docs/work_packages/wp-193d.md index 2c978ca..f4bafe5 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 none +> 🔗 Depends on [WP-193c](../work_packages/wp-193c.md) · Needed by [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-199.md b/docs/work_packages/wp-199.md index ad684d4..910bc20 100644 --- a/docs/work_packages/wp-199.md +++ b/docs/work_packages/wp-199.md @@ -2,9 +2,9 @@ > Engineering log · [Phase 17 — R1 pose-objective parity](../engineering/phase-17-pose-objective-parity.md) > -> 🔗 Depends on [WP-198](../work_packages/wp-198.md) · Needed by none +> 🔗 Depends on [WP-198](../work_packages/wp-198.md) · Needed by [WP-200](../work_packages/wp-200.md) > -> ⬅️ [WP-198](../work_packages/wp-198.md) · [Roadmap](../ROADMAP.md) +> ⬅️ [WP-198](../work_packages/wp-198.md) · [Roadmap](../ROADMAP.md) · [WP-200](../work_packages/wp-200.md) ➡️ diff --git a/docs/work_packages/wp-200.md b/docs/work_packages/wp-200.md new file mode 100644 index 0000000..054c5f7 --- /dev/null +++ b/docs/work_packages/wp-200.md @@ -0,0 +1,43 @@ +# WP-200 — replay native CUDA oriented networks + +> 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-199](../work_packages/wp-199.md) · Needed by none +> +> ⬅️ [WP-199](../work_packages/wp-199.md) · [Roadmap](../ROADMAP.md) + + + + + +## WP-200 — six native outputs, one recorded network + +**The oriented head is part of the recording.** The existing `capture_network` option now accepts `task: obb`. Backbone, neck and both head branches run their existing expressions through R43's graphed-callable API. The adapter retains the four dense detection fields and both angle tensors, rebuilding them by name. Assignment, rotated decoding, classification, ProbIoU, rotated L1, angle scoring, schedules, clipping, optimizer, EMA and validation stay in their current loop. No method arithmetic, gain, convention or frozen golden changes. + +**The complete native objective owns the gradient contract.** Its empty-target path retains zero gradients to both angle stems. An overridden objective can instead leave those gradients absent, so a changed training consumer, classification loss or oriented helper keeps the complete native call. Accumulation also stays native, as it does for pose: the existing shared recorder has no established accumulation contract. The ordinary network storage, hooks, precision, single-device CUDA and first-order constraints remain unchanged. One image signature is retained per fit; a temporary partial batch stays native and can return to the existing recording. The default stays off. + +**The stage was priced before implementation.** Four native CUDA steps occupied a 1.466 s instrumented window, with about 24.1% covered by recorded device activity. Network forward dispatch took 378.113 ms, including 174.956 ms in backbone and neck and 202.463 ms in the head. A separate complete fit with lightweight host timers took 68.263 s: network forwards took 15.943 s, or 23.35%, across 208 training and 32 validation calls. That includes validation, which remains native, and is an upper bound rather than a speed claim. The 21.261 s backward timer includes loss work as well as network work and is not wholly removable. Profiler runs are separate from the final timing comparisons. + +**The smaller recording did not clear the investigation's target.** A shared-backbone/neck-only screen preserved model, optimizer, scheduler, EMA and loss bytes, but changed complete-fit time from 68.277 to 63.800 s, a single-pair 6.56% reduction. It was set aside before the final cohort. The complete six-output screen preserved those same states and validation values and changed 68.742 to 59.853 s, a preliminary 12.93%; the repeated cohort, rather than that screen, determines the final claim. + +**Workload.** The registered nano oriented model trains on 104 public DOTA-v1.0 tiles with 2,075 annotations and validates on 13 tiles with 616 annotations. Training and validation use disjoint original source images, 16 and 4 respectively. Eight epochs, batch 4, 1024 px, BF16, seed 123, deterministic kernels, four persistent loader workers and one NVIDIA L4 on a g2-standard-8. Both arms enable prepared source caching, device HSV, epoch prefetch and the existing orthogonalization/update-preparation replay. The comparison uses separate upstream and candidate source trees in fresh processes. Timing brackets include model/trainer construction, cache and recording setup, augmentation, training, EMA, validation and checkpoint writes, with device synchronization at both boundaries. + +**[MEASURED] Five alternating fresh-process pairs.** One excluded warmup pair precedes the five measured pairs. Native and replay ranges do not overlap; every paired reduction is positive (11.669–13.572%). The exact model, optimizer, scheduler, EMA and per-step loss byte hashes agree in each pair and across native repeats, as do validation values. The package import path and hashes bind each arm to its separate source tree. The complete-fit median drops **12.011%**. + +| Metric, median (range across five runs) | Native | Replay | +| -- | -- | -- | +| Complete fit, seconds | 68.642 (68.213–69.333) | 60.397 (59.924–60.784) | +| Whole process, seconds | 73.795 (73.379–74.561) | 65.608 (65.169–65.987) | +| Warmed training step, ms | 183.275 (179.156–184.821) | 139.500 (138.453–141.266) | + +Warmed step diagnostics discard the first 64 steps and use CUDA events, synchronized before reading. Peak allocated memory is 2,572,018,688 bytes native and 1,889,886,208 bytes replay; peak reserved memory is 4,167,041,024 and 4,508,876,800 bytes, respectively, constant across the five runs. Extra graph-pool reservation is a tradeoff, not a capacity improvement. Dependency/compiler and OS caches are warm; each fit reconstructs its prepared image cache and CUDA recording. Package import and post-fit byte hashing are outside the fit timer and included in the separately reported process time. Runtime: Python 3.11.17, torch 2.9.1+cu129, Lightning 2.6.6, L4 driver 580.178.04. + +**The removed work is observed at its owner.** Separate four-step CUDA traces show 31,193 native `cudaLaunchKernel` calls and 12,609 with replay, while graph launches rise from 4 to 12, adding the network forward/backward recordings beside the existing optimizer replay. Device-event union is 353.374 and 358.045 ms; these instrumented windows support the dispatch diagnosis, not a kernel-speed claim. + +**A larger-batch probe limits the claim.** One batch-8 pair, with the same eight epochs and 1024 px inputs, changes complete-fit time from 60.494 to 58.014 s, only 4.099%. State and validation parity still hold, but that one pair is below the investigation's 8% target and is not a repeated performance claim. The five-pair headline remains scoped to batch 4. + +**Validation.** Thirty oriented/constructor cases passed on both torch 2.9.1 and 2.4.0; the floor retains its existing BF16 eager fallback, while FP32 and FP16 exercise real replay. A silent eager reversion and an unsafe objective-guard bypass each make the committed Trainer tests fail. Five additional helper-replacement cases pass on CPU after review. Local CPU checks and source doctests, with those cases included, report 17 passed and 19 CUDA skips. Strict mypy and the full pre-commit hooks pass. No frozen golden is changed. + +**Coverage and limits.** Focused tests compare all six fields, actual CUDA launches, empty-target gradient presence, temporary shape fallback, gradient clipping, accumulation retirement (native fallback), FP32/FP16/BF16, scaler overflow and recovery, native-objective fallback, EMA validation and fit cleanup. FP16 fixtures use a unit initial scale plus one injected overflow so at least one real optimizer update must occur. A silent native reversion must fail the launch assertions; bypassing the objective guard must be rejected by the absent-angle-gradient case. Full `make gate`, accelerator goldens and long-run acceptance accuracy are outside this validation run. + + diff --git a/scripts/_tests/test_audit_docs_present.py b/scripts/_tests/test_audit_docs_present.py index a9754c6..a895b11 100644 --- a/scripts/_tests/test_audit_docs_present.py +++ b/scripts/_tests/test_audit_docs_present.py @@ -83,7 +83,8 @@ def _write_required_docs(docs_dir: Path, repo_root: Path) -> None: _write(docs_dir / "ESCALATION.md", "escalation\n") _write( docs_dir / "ROADMAP.md", - "199 numbered work packages.\n" + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, 200)), + f"{audit._WP_FLOOR} numbered work packages.\n" + + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, audit._WP_FLOOR + 1)), ) _write(docs_dir / "DATASETS.md", "datasets\n") _write(docs_dir / "TRAINING.md", "training\n") @@ -305,11 +306,16 @@ class TestCheckRoadmapWpIdsUniqueAndComplete: def test_flags_floor(self, tmp_path: Path) -> None: """A roadmap below the work-package floor is reported even when contiguous.""" _write_required_docs(tmp_path / "docs", tmp_path) - _write(tmp_path / "docs" / "ROADMAP.md", "1 numbered work packages.\n| 001 | ... |\n") + floor = audit._WP_FLOOR + _write( + tmp_path / "docs" / "ROADMAP.md", + f"{floor - 1} numbered work packages.\n" + + "".join(f"| {i:03d} | a | b | c | d | ✅ |\n" for i in range(1, floor)), + ) violations = audit.check_roadmap_wp_ids_unique_and_complete(tmp_path / "docs", tmp_path) - assert any("shrank below 199" in violation for violation in violations) + assert violations == [f"roadmap shrank below {floor} work packages: {floor - 1}"] def test_clean_at_floor(self, tmp_path: Path) -> None: """A contiguous roadmap at the floor count reports no violation.""" @@ -365,12 +371,13 @@ def test_flags_mismatched_count(self, tmp_path: Path) -> None: _write_required_docs(tmp_path / "docs", tmp_path) text = (tmp_path / "docs" / "ROADMAP.md").read_text(encoding="utf-8") _write( - tmp_path / "docs" / "ROADMAP.md", text.replace("199 numbered work packages.", "5 numbered work packages.") + tmp_path / "docs" / "ROADMAP.md", + text.replace(f"{audit._WP_FLOOR} numbered work packages.", "5 numbered work packages."), ) violations = audit.check_roadmap_header_count(tmp_path / "docs", tmp_path) - assert violations == ["header states 5 numbered work packages but the table carries 199"] + assert violations == [f"header states 5 numbered work packages but the table carries {audit._WP_FLOOR}"] def test_clean_when_counts_agree(self, tmp_path: Path) -> None: """A stated package count matching the numbered rows reports no violation.""" diff --git a/scripts/lint/audit_docs_present.py b/scripts/lint/audit_docs_present.py index b950d46..8301793 100644 --- a/scripts/lint/audit_docs_present.py +++ b/scripts/lint/audit_docs_present.py @@ -31,7 +31,7 @@ #: Lowest work-package count ROADMAP.md is allowed to hold. A ratchet, not a target: #: contiguity alone would not notice the last row being deleted. Raise it when adding #: a work package; never lower it. -_WP_FLOOR = 199 +_WP_FLOOR = 200 #: Lowest decision count DECISIONS.md is allowed to hold. A ratchet, not a target: #: contiguity alone would not notice the last row being deleted, since what remains diff --git a/src/lucid_yolo/ptl/loss_capture.py b/src/lucid_yolo/ptl/loss_capture.py index 3e5019f..794cbf2 100644 --- a/src/lucid_yolo/ptl/loss_capture.py +++ b/src/lucid_yolo/ptl/loss_capture.py @@ -91,17 +91,28 @@ def _functions(source: DualBranchLoss) -> tuple[object, ...]: _EXPECTED = _functions(DualBranchLoss()) -def _source_signature(source: DualBranchLoss) -> tuple[object, ...] | None: - """Certify the built-in objects and scalar settings, excluding the live alpha input.""" +def _native_loss(source: DualBranchLoss) -> bool: + """Tell whether the built-in loss objects and the implementations they call are unchanged. + + Only identity is judged: exact component types and the callables :func:`_functions` + reads. Scalar settings are not, so a caller whose loss runs outside any recording can + rely on this without inheriting the workspace limits :func:`_source_signature` adds. + """ if type(source) is not DualBranchLoss: - return None + return False objects = (source, source.o2m_assigner, source.o2o_assigner, source._o2m_loss, source._o2o_loss) - if ( + return ( tuple(type(obj) for obj in objects) - != (DualBranchLoss, SmallTargetAssigner, UniqueAssigner, DetectionBranchLoss, DetectionBranchLoss) - or _functions(source) != _EXPECTED - ): + == (DualBranchLoss, SmallTargetAssigner, UniqueAssigner, DetectionBranchLoss, DetectionBranchLoss) + and _functions(source) == _EXPECTED + ) + + +def _source_signature(source: DualBranchLoss) -> tuple[object, ...] | None: + """Certify the built-in objects and scalar settings, excluding the live alpha input.""" + if not _native_loss(source): return None + objects = (source, source.o2m_assigner, source.o2o_assigner, source._o2m_loss, source._o2o_loss) fields = ("topk", "alpha", "beta", "eps", "s_min", "s_ref") settings = tuple(getattr(assigner, name) for assigner in objects[1:3] for name in fields) settings += tuple(getattr(branch, name) for branch in objects[3:] for name in ("box_gain", "cls_gain", "l1_gain")) diff --git a/src/lucid_yolo/ptl/module.py b/src/lucid_yolo/ptl/module.py index 186f97e..0720a9c 100644 --- a/src/lucid_yolo/ptl/module.py +++ b/src/lucid_yolo/ptl/module.py @@ -189,6 +189,7 @@ ) from lucid_yolo.eval.dota_eval import MAX_DETECTIONS, evaluate_rotated_map, rotated_detections_to_predictions from lucid_yolo.eval.segment_decode import decode_instance_masks +from lucid_yolo.losses import oriented_loss from lucid_yolo.losses.dual_loss import DualBranchLoss, DualLossOutput from lucid_yolo.losses.keypoint_nll_loss import LaplaceNLLLoss from lucid_yolo.losses.mask_loss import instance_mask_loss @@ -216,7 +217,7 @@ from lucid_yolo.optim.musgd import MuSGD from lucid_yolo.optim.schedule import warmup_decay_factor from lucid_yolo.ptl.coco_backend import build_mean_average_precision -from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle +from lucid_yolo.ptl.loss_capture import _LossCapture, _LossCaptureLifecycle, _native_loss from lucid_yolo.ptl.network_capture import ( _detection_stages, _module_breach, @@ -791,21 +792,21 @@ def _validate_replay_options( ) -> None: """Reject replay options outside the objectives they record. - Network replay records the detection or pose 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 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 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 or - keypoints, 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 + 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). """ - if capture_network and (task not in ("detect", "keypoints") or compile_step): - raise ValueError("capture_network=True requires task='detect' or task='keypoints', and compile_step=False") + 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_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): @@ -815,6 +816,37 @@ def _validate_replay_options( ) +def _oriented_helpers() -> tuple[object, ...]: + """Read the decode and oriented-loss helpers the native objective calls by global name. + + Resolved at call time, so the import-time call that builds :data:`_ORIENTED_HELPERS` + keeps the originals and every later call sees the current bindings. + """ + return ( + decode_ltrb, + decode_rboxes, + oriented_branch_terms, + vars(oriented_loss)["square_angle_loss"], + oriented_loss._no_ground_truth, + ) + + +def _oriented_replay_breach(module: DetectionLitModule) -> str | None: + """Keep full-network replay on the native objective that connects all six outputs. + + The native empty-target path retains zero gradients for the angle stems. An + overridden objective may leave them absent instead, which a recorded backward + would materialize. Changed consumers therefore keep their complete native call. + """ + for name, expected in _ORIENTED_METHODS: + method = getattr(module, name) + if getattr(method, "__func__", method) is not expected: + return "the native oriented objective is overridden" + if _oriented_helpers() != _ORIENTED_HELPERS: + return "a native oriented objective helper changed" + return None if _native_loss(module.loss) else "the native loss or an assignment implementation changed" + + class DetectionLitModule(LightningModule): """Lightning detector: model stack, dual-branch loss, and MuSGD (WP-034). @@ -973,10 +1005,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 or keypoint network stages' CUDA forward and backward + 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"`` or ``task="keypoints"``, ``compile_step=False`` and + ``task="detect"``, ``task="keypoints"`` or ``task="obb"``, ``compile_step=False`` and single-device CUDA automatic optimization. Pose records backbone and neck; its head stays native to preserve absent gradients in batches with no assigned positives. Pose training stays eager @@ -985,7 +1017,10 @@ class DetectionLitModule(LightningModule): module is hooked, overridden or compiled. The trainable set is fixed per fit at the first recording: a fully frozen start, or a later freeze or unfreeze of a recorded parameter, keeps the rest of that fit eager after one warning. - Either task stays eager when the LightningModule carries a hook, overrides + 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. Assignment, losses, schedules, optimizer, EMA and evaluation stay eager. One image signature is recorded per fit; partial batches stay eager and @@ -1014,9 +1049,9 @@ 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 or keypoints, - ``capture_loss`` with a non-detection task, ``capture_segment_heads`` with a - non-segmentation task, or any of them with ``compile_step=True`` (for + 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`` + with a non-segmentation task, or any of them with ``compile_step=True`` (for ``capture_segment_heads``, forward compatibility with a compiled segment objective), if ``keypoint_loss`` is not a known keypoint objective, if ``task="keypoints"`` was asked for without a ``num_keypoints``, or if a @@ -1494,10 +1529,10 @@ def _detect_objective( def _training_network(self, images: Tensor) -> DualHeadOutput: """Run the dual head through the training recording when its contract holds. - Detection records all three stages; pose records only backbone and neck. - Its optional point stems stay native so batches with no assigned positives + Detection and oriented detection record all three stages; pose records only backbone and neck. + Pose's optional point stems stay native so batches with no assigned positives keep their absent gradients, rather than zero-valued captured gradients. - Both tasks keep changing instance counts and loss arithmetic outside capture. + All tasks keep changing instance counts and loss arithmetic outside capture. """ head_out: DualHeadOutput if ( @@ -1511,8 +1546,9 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: ): # 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 == "keypoints" and self.trainer.accumulate_grad_batches != 1: - self._network_capture.retire("keypoint replay requires accumulate_grad_batches=1") + 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") if self._task == "keypoints": breach = _pose_head_breach(self.head, self._pose_head_topology) if breach is not None: @@ -1533,6 +1569,8 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: elif self._task == "keypoints": features = self._network_capture.run_features(images, (self.backbone, self.neck)) head_out = self(images) if features is None else cast("DualHeadOutput", self.head(features)) + elif self._task == "obb": + head_out = self._training_oriented_network(images) else: eager = cast("Callable[[Tensor], DualHeadOutput]", self) head_out = self._network_capture.run(images, (self.backbone, self.neck, self.head), eager) @@ -1540,6 +1578,15 @@ def _training_network(self, images: Tensor) -> DualHeadOutput: head_out = self(images) return head_out + def _training_oriented_network(self, images: Tensor) -> DualHeadOutput: + """Record the complete oriented network only while its native consumers stay intact.""" + reason = _oriented_replay_breach(self) + if reason is not None: + self._network_capture.retire(reason) + return cast("DualHeadOutput", self(images)) + eager = cast("Callable[[Tensor], DualHeadOutput]", self) + return self._network_capture.run_oriented(images, (self.backbone, self.neck, self.head), eager) + def _alpha_tensor(self, device: torch.device) -> Tensor: """Lift the current :attr:`alpha` to a 0-d ``float64`` tensor on ``device`` for the compiled region. @@ -2795,3 +2842,11 @@ 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 + +#: Consumers whose native calls retain gradients to every oriented output, even on empty batches. +_ORIENTED_METHODS = tuple( + (name, getattr(DetectionLitModule, name)) + for name in ("training_step", "_shared_step", "_task_extra_loss", "_oriented_extra_loss", "_branch_oriented_terms") +) +#: The oriented objective helpers as imported, against which :func:`_oriented_replay_breach` compares. +_ORIENTED_HELPERS = _oriented_helpers() diff --git a/src/lucid_yolo/ptl/network_capture.py b/src/lucid_yolo/ptl/network_capture.py index f5fbe72..ce3bf52 100644 --- a/src/lucid_yolo/ptl/network_capture.py +++ b/src/lucid_yolo/ptl/network_capture.py @@ -2,8 +2,8 @@ """Opt-in replay of fixed network stages' CUDA forward and backward work. The network has a fixed image signature while its ground-truth count changes per -batch. Detection records backbone, neck and head; pose records backbone and neck. -Both leave assignment, losses, schedules, +batch. Detection and oriented detection record backbone, neck and head; pose records backbone and neck. +All three 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. @@ -72,6 +72,53 @@ def rebuild(tensors: tuple[Tensor, ...]) -> DualHeadOutput: return DualHeadOutput(o2m_cls=o2m_cls, o2m_box=o2m_box, o2o_cls=o2o_cls, o2o_box=o2o_box) +class _OrientedNetwork(nn.Module): + """Expose both native angle tensors beside the four dense detection fields. + + Examples: + ```pycon + >>> from lucid_yolo.ptl import DetectionLitModule + >>> model = DetectionLitModule(0.34, 0.25, 64, 2, task="obb").eval() + >>> network = _OrientedNetwork(model.backbone, model.neck, model.head) + >>> with torch.no_grad(): + ... tensors = network(torch.zeros(2, 3, 64, 64)) + >>> len(tensors), network.rebuild(tensors).o2m_angle.shape + (6, torch.Size([2, 84, 1])) + + ``` + """ + + def __init__(self, backbone: nn.Module, neck: nn.Module, head: nn.Module) -> None: + super().__init__() + self.backbone, self.neck, self.head = backbone, neck, head + + def forward(self, images: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Run the existing stages and retain the six tensors the native oriented objective reads.""" + output = _detection_stages(self.backbone, self.neck, self.head, images) + assert output.o2m_angle is not None and output.o2o_angle is not None + return ( + output.o2m_cls, + output.o2m_box, + output.o2o_cls, + output.o2o_box, + output.o2m_angle, + output.o2o_angle, + ) + + @staticmethod + def rebuild(tensors: tuple[Tensor, ...]) -> DualHeadOutput: + """Rebuild every recorded field by name, leaving absent modalities absent.""" + o2m_cls, o2m_box, o2o_cls, o2o_box, o2m_angle, o2o_angle = tensors + return DualHeadOutput( + o2m_cls=o2m_cls, + o2m_box=o2m_box, + o2o_cls=o2o_cls, + o2o_box=o2o_box, + o2m_angle=o2m_angle, + o2o_angle=o2o_angle, + ) + + class _FeatureNetwork(nn.Module): """Expose backbone and neck without fixing which optional head gradients exist. @@ -321,7 +368,7 @@ def retire(self, reason: str, *, stacklevel: int = 3) -> None: stacklevel: The frame the warning names, counted as :func:`warnings.warn` counts from this method. The default names the caller of whoever called ``retire``; :meth:`_run` passes 4 so a retirement it decides still names - the code that called :meth:`run` or :meth:`run_features`. + the code that called :meth:`run`, :meth:`run_features` or :meth:`run_oriented`. """ if self._disabled: return @@ -363,6 +410,15 @@ def eager(_: Tensor) -> _Features | None: rebuild: Callable[[tuple[Tensor, ...]], _Features | None] = _FeatureNetwork.rebuild return self._run(images, stages, eager, lambda: _FeatureNetwork(*stages), rebuild) + def run_oriented( + self, + images: Tensor, + stages: tuple[nn.Module, nn.Module, nn.Module], + eager: Callable[[Tensor], DualHeadOutput], + ) -> DualHeadOutput: + """Replay all six network outputs consumed by the unchanged native oriented objective.""" + return self._run(images, stages, eager, lambda: _OrientedNetwork(*stages), _OrientedNetwork.rebuild) + def _run( self, images: Tensor, @@ -378,8 +434,8 @@ def _run( stages: The caller's stage modules, whose identities the recording keeps. eager: Produces the value returned whenever this call does not replay. It is a fallback value, not necessarily the native result: :meth:`run` passes the - native forward, while :meth:`run_features` returns ``None`` so its caller - makes the full native call itself. + native forward, and so does :meth:`run_oriented`, while :meth:`run_features` + returns ``None`` so its caller makes the full native call itself. build: Composes the stages into the module checked for eligibility until a recording exists, and captured by the first one. rebuild: Turns the replayed output tensors back into the caller's value. diff --git a/tests/ptl/test_network_capture.py b/tests/ptl/test_network_capture.py index 18c4194..fd99804 100644 --- a/tests/ptl/test_network_capture.py +++ b/tests/ptl/test_network_capture.py @@ -11,6 +11,7 @@ import copy import gc import pickle +import re import warnings import weakref from collections.abc import Callable, Iterator, Sequence @@ -90,14 +91,19 @@ def test_default_has_no_callback_or_registered_state() -> None: "kwargs", [ pytest.param({"task": "segment"}, id="segment"), - pytest.param({"task": "obb"}, id="obb"), + pytest.param({"task": "obb", "compile_step": True}, id="obb-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.""" - with pytest.raises(ValueError, match="requires task='detect' or task='keypoints', and compile_step=False"): + """Unsupported tasks and the separate compiler path fail 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" + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _module(capture_network=True, **kwargs) diff --git a/tests/ptl/test_obb_training.py b/tests/ptl/test_obb_training.py index 01ff7d0..dedb0a4 100644 --- a/tests/ptl/test_obb_training.py +++ b/tests/ptl/test_obb_training.py @@ -399,17 +399,23 @@ def test_image_without_instances_gives_finite_loss() -> None: With no ground truth the padded rotated axis has length zero, which no ``gather`` can index; the term has to answer with an exact zero that is still connected to - the predictions. + the predictions. Connected means the backward reaches every angle stem of both + branches with a gradient, the native premise full-network replay rests on: an empty + batch that left an angle gradient absent would make a recorded backward differ. """ module = _tiny_module() recorder = _LogRecorder() module.log = recorder # type: ignore[method-assign] loss = module.training_step((torch.randn(1, 3, _IMG_SIZE, _IMG_SIZE), [Targets.empty()]), 0) + loss.backward() assert bool(torch.isfinite(loss)) assert recorder.values["train/rbox"] == 0.0 assert recorder.values["train/angle"] == 0.0 + angles = {name: parameter for name, parameter in module.named_parameters() if "angle_stems" in name} + assert any(name.startswith("head.o2o.") for name in angles) and any(name.startswith("head.o2m.") for name in angles) + assert all(parameter.grad is not None for parameter in angles.values()), angles.keys() def test_rotated_candidacy_reaches_assigner() -> None: diff --git a/tests/ptl/test_oriented_network_capture.py b/tests/ptl/test_oriented_network_capture.py new file mode 100644 index 0000000..43c8fcc --- /dev/null +++ b/tests/ptl/test_oriented_network_capture.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Oriented replay retains angle fields, native gradients and validation transitions.""" + +from __future__ import annotations + +import copy +import warnings +from collections.abc import Callable, Iterator +from dataclasses import fields +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +import torch +from pytorch_lightning import Callback, Trainer +from pytorch_lightning.plugins.precision import MixedPrecision +from torch.utils.data import DataLoader + +from lucid_yolo.cli.train import DetectionCLI, packaged_config +from lucid_yolo.data.rotated_geom import rboxes_to_polygons +from lucid_yolo.data.targets import Targets +from lucid_yolo.models.heads.detect import DualHeadOutput +from lucid_yolo.ptl import DetectionLitModule, EMACallback +from lucid_yolo.ptl.datamodule import DetectionDataModule +from lucid_yolo.ptl.loss_capture import _source_signature +from lucid_yolo.ptl.module import _oriented_replay_breach +from lucid_yolo.ptl.network_capture import _OrientedNetwork + + +@pytest.fixture(autouse=True) +def _deterministic_kernels() -> Iterator[None]: + """Keep native comparisons deterministic and restore the caller's policy afterwards.""" + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True) + try: + yield + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + + +def _oriented(**kwargs: Any) -> DetectionLitModule: + """Build a small ordinary oriented model whose optimizer fits in a focused test. + + Examples: + >>> _oriented().task + 'obb' + """ + options = {"lr": 1e-4, "warmup_epochs": 0.0, "lrf": 1.0, "w_muon": 0.0, "w_sgd": 1.0} + options.update(kwargs) + return DetectionLitModule(0.34, 0.25, 64, 2, task="obb", **options) + + +def _same(left: Any, right: Any) -> None: + """Compare tensor bytes and recursive training state, including signed zero. + + Examples: + >>> _same({"counter": torch.tensor(2)}, {"counter": torch.tensor(2)}) + """ + if isinstance(left, torch.Tensor): + assert left.shape == right.shape and left.dtype == right.dtype + assert torch.equal( + left.detach().contiguous().reshape(-1).view(torch.uint8), + right.detach().contiguous().reshape(-1).view(torch.uint8), + ) + elif isinstance(left, dict): + assert left.keys() == right.keys() + for key, value in left.items(): + _same(value, right[key]) + elif isinstance(left, (list, tuple)): + assert len(left) == len(right) + for first, second in zip(left, right, strict=True): + _same(first, second) + else: + assert left == right + + +def _batches(layout: str) -> list[tuple[torch.Tensor, list[Targets]]]: + """Keep full, empty and temporary partial batches in one native loader contract. + + Examples: + >>> [len(targets) for _, targets in _batches("partial")] + [2, 1, 2] + """ + image = torch.rand(2, 3, 64, 64) + rboxes = torch.tensor([[32.0, 32.0, 32.0, 12.0, 0.3]]) + corners = rboxes_to_polygons(rboxes) + target = Targets( + boxes=torch.cat((corners.amin(dim=1), corners.amax(dim=1)), dim=1), + labels=torch.tensor([1]), + rboxes=rboxes, + ) + empty = Targets(boxes=torch.empty(0, 4), labels=torch.empty(0, dtype=torch.int64), rboxes=torch.empty(0, 5)) + batches = [(image.clone(), [copy.deepcopy(target), copy.deepcopy(target)]) for _ in range(3)] + if layout == "empty": + batches[1] = (image.clone(), [copy.deepcopy(empty), copy.deepcopy(empty)]) + elif layout == "partial": + batches[1] = (image[:1].clone(), [copy.deepcopy(target)]) + return batches + + +class _ReplayProbe(Callback): + """Observe real graph launches and angle-gradient presence after every training batch.""" + + def __init__(self, launches: list[int]) -> None: + self.launches = launches + self.deltas: list[int] = [] + self.angle_gradients: list[bool] = [] + self.losses: list[torch.Tensor] = [] + self.start = 0 + + def on_train_batch_start(self, *args: Any, **kwargs: Any) -> None: + """Mark the real replay count immediately before this batch's forward.""" + self.start = len(self.launches) + + def on_train_batch_end(self, trainer: Trainer, module: DetectionLitModule, outputs: Any, *args: Any) -> None: + """Retain launched work, live gradients and loss bytes after backward and the optimizer.""" + self.deltas.append(len(self.launches) - self.start) + angles = [parameter for name, parameter in module.named_parameters() if "angle_stems" in name] + assert angles + self.angle_gradients.append(all(parameter.grad is not None for parameter in angles)) + self.losses.append(outputs["loss"].detach().clone()) + + +class _OverflowFirst(Callback): + """Make one backward non-finite, then require the scaler to resume real updates.""" + + def __init__(self) -> None: + self.pending = True + + def on_before_backward(self, trainer: Trainer, module: DetectionLitModule, loss: torch.Tensor) -> None: + """Inject through the public callback without changing which network outputs are consumed.""" + if self.pending: + loss.mul_(float("inf")) + self.pending = False + + +def test_oriented_adapter_preserves_fields() -> None: + """The six-field adapter cannot swap the two angle branches or lose optional fields.""" + model = _oriented().eval() + with torch.no_grad(): + assert model.head.o2o.angle_stems is not None + for parameter in model.head.o2o.angle_stems.parameters(): + parameter.add_(0.1) + image = torch.rand(2, 3, 64, 64) + expected = model(image) + network = _OrientedNetwork(model.backbone, model.neck, model.head) + actual = network.rebuild(network(image)) + assert not torch.equal(expected.o2m_angle, expected.o2o_angle) + for field in fields(DualHeadOutput): + _same(getattr(expected, field.name), getattr(actual, field.name)) + + +@pytest.mark.parametrize("stage", ["fit", "validate", "test", "predict"]) +def test_oriented_default_and_cpu_fallback(stage: str) -> None: + """Enabling the existing flag changes no checkpoint keys and stays native without CUDA. + + Every Lightning stage sets the module up, and only ``fit`` trains, so only ``fit`` + names the ignored flag: an evaluation run reporting a replay it never would use is + noise, and a CPU setup that turned recording on would be a silent regression. The + recorder entry point raises if anything reaches it, so the CPU setup is proven inert. + """ + native, replayed = _oriented(), _oriented(capture_network=True) + assert native.state_dict().keys() == replayed.state_dict().keys() + assert not native.hparams.capture_network and replayed.hparams.capture_network + assert _oriented_replay_breach(replayed) is None + replayed.trainer = Trainer(accelerator="cpu", logger=False, enable_checkpointing=False) + with ( + patch("torch.cuda.make_graphed_callables", side_effect=AssertionError("CPU must stay eager")), + warnings.catch_warnings(record=True) as caught, + ): + warnings.simplefilter("always") + replayed.setup(stage) + ignored = [str(item.message) for item in caught if str(item.message).startswith("capture_network=True is ignored")] + assert len(ignored) == int(stage == "fit"), ignored + assert not replayed._capture_network_active + + +def test_oriented_cli_accepts_network_replay() -> None: + """The packaged oriented recipe builds the flag, keeps rotated-target transport and passes the replay guard. + + A shipped setting that tripped the guard would leave every user of the recipe on the + eager path with only a retirement warning to show for it. + """ + cli = DetectionCLI( + DetectionLitModule, + DetectionDataModule, + args=[ + "--config", + str(packaged_config("obb_nano_smoke")), + "--model.capture_network", + "true", + "--trainer.accelerator", + "cpu", + "--trainer.logger", + "false", + ], + run=False, + ) + assert cli.model.task == "obb" and cli.model.hparams.capture_network + assert cli.datamodule._rotated_targets + assert _oriented_replay_breach(cli.model) is None + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA graph replay") +@pytest.mark.parametrize("precision", ["32-true", "16-mixed", "bf16-mixed"]) +@pytest.mark.parametrize("layout", ["uniform", "empty", "partial"]) +@pytest.mark.parametrize("accumulation", [1, 2]) +def test_oriented_trainer_preserves_native_state( + precision: str, + layout: str, + accumulation: int, + graph_recordings: MagicMock, + bf16_replay_expected: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Actual replay survives EMA validation and preserves empty-target gradients and temporary fallback. + + A replaying arm that never records names exactly one cause, and the cause is pinned: + accumulation is checked before the recorder reads the precision policy, so with both + breached the accumulation reason is the one reported. A retirement for an unrelated + cause would otherwise satisfy the count while hiding the regression behind it. + """ + torch.manual_seed(42) + batches = _batches(layout) + retirements = ( + ["accumulate_grad_batches=1"] + if accumulation == 2 + else ["BF16 replay requires"] + if precision == "bf16-mixed" and not bf16_replay_expected + else [] + ) + initial = _oriented().state_dict() + real_replay = torch.cuda.CUDAGraph.replay + launches: list[int] = [] + + def observe(graph: torch.cuda.CUDAGraph) -> None: + """Count actual graph execution while retaining the real runtime call.""" + launches.append(id(graph)) + real_replay(graph) + + monkeypatch.setattr(torch.cuda.CUDAGraph, "replay", observe) + results = [] + enabled = torch.are_deterministic_algorithms_enabled() + warn_only = torch.is_deterministic_algorithms_warn_only_enabled() + try: + for capture in (False, True): + torch.manual_seed(42) + # A decaying LR keeps a real scheduler on both arms, so its state joins the oracle. + model = _oriented(capture_network=capture, lrf=0.1) + model.load_state_dict(initial) + probe, ema = _ReplayProbe(launches), EMACallback() + callbacks: list[Callback] = [probe, ema] + precision_options: dict[str, Any] = {"precision": precision} + if precision == "16-mixed": + # The native short fixture skips every update at the default scale. + # Unit scale plus one explicit overflow proves both skipping and recovery. + scaler = torch.amp.GradScaler("cuda", init_scale=1.0) + precision_options = {"plugins": [MixedPrecision("16-mixed", "cuda", scaler=scaler)]} + callbacks.append(_OverflowFirst()) + trainer = Trainer( + accelerator="cuda", + devices=1, + max_epochs=2, + gradient_clip_val=1.0, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + num_sanity_val_steps=0, + deterministic=True, + callbacks=callbacks, + accumulate_grad_batches=accumulation, + **precision_options, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer.fit( + model, + train_dataloaders=DataLoader(copy.deepcopy(batches), batch_size=None), + val_dataloaders=DataLoader(copy.deepcopy(batches[:1]), batch_size=None), + ) + recording = capture and accumulation == 1 and (precision != "bf16-mixed" or bf16_replay_expected) + retired = [ + str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays") + ] + assert len(retired) == int(capture and not recording), retired + expected = retirements if capture else [] + assert all(fragment in message for fragment, message in zip(expected, retired, strict=True)), retired + assert graph_recordings.call_count == int(recording) + assert len(probe.deltas) == 6 and all(probe.angle_gradients) + assert trainer.global_step == (3 + accumulation - 1) // accumulation * 2 + for index, delta in enumerate(probe.deltas): + active = recording and not (layout == "partial" and index % 3 == 1) + assert delta >= 2 if active else delta == 0 + assert trainer.optimizers[0].state + assert model._network_capture._network is None + scaler = getattr(trainer.precision_plugin, "scaler", None) + if scaler is not None: + assert scaler.get_scale() < 1.0 + results.append( + copy.deepcopy( + ( + model.state_dict(), + trainer.optimizers[0].state_dict(), + trainer.lr_scheduler_configs[0].scheduler.state_dict(), + ema.state_dict(), + probe.losses, + None if scaler is None else scaler.state_dict(), + trainer.callback_metrics, + ) + ) + ) + _same(*results) + finally: + torch.use_deterministic_algorithms(enabled, warn_only=warn_only) + + +@pytest.mark.parametrize( + "method", ["training_step", "_shared_step", "_task_extra_loss", "_oriented_extra_loss", "_branch_oriented_terms"] +) +def test_oriented_objective_override_stays_native(method: str, monkeypatch: pytest.MonkeyPatch) -> None: + """A changed consumer cannot silently replace absent gradients with recorded zeros.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(model, method, lambda *args: torch.tensor(0.0)) + assert _oriented_replay_breach(model) == "the native oriented objective is overridden" + + +def test_oriented_classification_override_stays_native(monkeypatch: pytest.MonkeyPatch) -> None: + """An unchanged model does not certify a replaced classification consumer.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(model.loss._o2m_loss, "_classification_loss", lambda *args: torch.tensor(0.0)) + assert _oriented_replay_breach(model) == "the native loss or an assignment implementation changed" + + +def test_oriented_loss_setting_change_keeps_replay(monkeypatch: pytest.MonkeyPatch) -> None: + """A scalar loss setting outside the network recording leaves oriented replay eligible. + + The oriented loss always runs eagerly after the recorded network, so an assigner epsilon + the detection loss recording refuses cannot change which outputs carry gradients. Both + halves are asserted: the loss certificate still refuses the value, and the network + guard does not inherit that refusal. + """ + model = _oriented(capture_network=True) + monkeypatch.setattr(model.loss.o2m_assigner, "eps", 1e-40) + assert _source_signature(model.loss) is None + assert _oriented_replay_breach(model) is None + + +@pytest.mark.parametrize( + "target", + [ + pytest.param("lucid_yolo.ptl.module.decode_ltrb", id="decode-ltrb"), + pytest.param("lucid_yolo.ptl.module.decode_rboxes", id="decode-rboxes"), + pytest.param("lucid_yolo.ptl.module.oriented_branch_terms", id="oriented-branch-terms"), + pytest.param("lucid_yolo.losses.oriented_loss.square_angle_loss", id="square-angle-loss"), + pytest.param("lucid_yolo.losses.oriented_loss._no_ground_truth", id="no-ground-truth"), + ], +) +def test_oriented_helper_replacement_stays_native(target: str, monkeypatch: pytest.MonkeyPatch) -> None: + """A swapped decode or loss helper retires replay, the oriented twin of WP-199's delegated-stage check.""" + model = _oriented(capture_network=True) + monkeypatch.setattr(target, lambda *args, **kwargs: None) + assert _oriented_replay_breach(model) == "a native oriented objective helper changed" + + +class _CudaFlagged(torch.Tensor): + """CPU tensor that reports CUDA residency so the CUDA-gated oriented routing is reachable offline. + + Operations unwrap to plain tensors, so the real model still runs on the CPU. A copy of + the pose suite's class: ``tests/`` is not a package, so one test module cannot import + another's helpers. + """ + + is_cuda = True + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): # type: ignore[no-untyped-def] + """Run the operation on the underlying data and return an ordinary tensor.""" + with torch._C.DisableTorchFunctionSubclass(): + return func(*args, **(kwargs or {})) + + +def _routed_oriented(accumulation: int = 1) -> DetectionLitModule: + """Build an oriented module that satisfies every training-time gate of the full-network route. + + Examples: + >>> _routed_oriented(2).trainer.accumulate_grad_batches + 2 + """ + model = _oriented(capture_network=True) + model._trainer = SimpleNamespace(training=True, accumulate_grad_batches=accumulation) + model._capture_network_active = True + return model + + +def _cuda_flagged_images() -> torch.Tensor: + """Images that pass the ``is_cuda`` gate while living on the CPU. + + Examples: + >>> _cuda_flagged_images().is_cuda + True + """ + return torch.rand(2, 3, 64, 64).as_subclass(_CudaFlagged) + + +def _override_objective(model: DetectionLitModule) -> None: + """Replace one native oriented consumer on this instance, as a user customization would. + + Examples: + >>> model = _oriented() + >>> _override_objective(model) + >>> "_task_extra_loss" in vars(model) + True + """ + model._task_extra_loss = lambda *args: torch.tensor(0.0) # type: ignore[method-assign] + + +def test_oriented_routing_dispatches_full_network_to_recorder() -> None: + """An unbroken oriented module hands all three stages to the oriented recorder and returns its output. + + The CUDA-only route cannot execute on CPU runners, so images spoof ``is_cuda`` and the + recorder is replaced: what is pinned is the dispatch itself, the oriented entry point + with the complete stage tuple and the module as its eager fallback, never the + detection route, and no retirement warning. + """ + model = _routed_oriented() + images = _cuda_flagged_images() + replayed = object() + with ( + patch.object(model._network_capture, "run_oriented", return_value=replayed) as oriented, + patch.object(model._network_capture, "run") as detect, + warnings.catch_warnings(), + ): + warnings.simplefilter("error") + output = model._training_network(images) + oriented.assert_called_once_with(images, (model.backbone, model.neck, model.head), model) + detect.assert_not_called() + assert output is replayed + + +def test_oriented_routing_leaves_cpu_images_native() -> None: + """Images that are not on CUDA never reach the oriented recorder, however many gates pass. + + Every other training-time gate is satisfied, so the image residency check alone + keeps the route native: the recorder is never called, no retirement is warned, and + all six fields equal the module's own eager forward byte for byte. + """ + model = _routed_oriented() + images = torch.rand(2, 3, 64, 64) + expected = model(images) + with ( + patch.object(model._network_capture, "run_oriented") as oriented, + warnings.catch_warnings(), + ): + warnings.simplefilter("error") + output = model._training_network(images) + oriented.assert_not_called() + _same( + [getattr(expected, field.name) for field in fields(DualHeadOutput)], + [getattr(output, field.name) for field in fields(DualHeadOutput)], + ) + + +@pytest.mark.parametrize( + ("accumulation", "mutate", "recorder_calls", "reason"), + [ + pytest.param(2, lambda model: None, 1, "oriented replay requires accumulate_grad_batches=1", id="accumulation"), + pytest.param(1, _override_objective, 0, "the native oriented objective is overridden", id="objective-override"), + ], +) +def test_oriented_routing_retires_then_runs_native( + accumulation: int, mutate: Callable[[DetectionLitModule], None], recorder_calls: int, reason: str +) -> None: + """A broken oriented contract retires the recording once and returns the full native call. + + Accumulation retires before dispatch, so the retired recorder is still reached and hands + back its eager fallback; an overridden consumer retires before the recorder is reached + at all. Either way the output is the module's own forward, run exactly once. + """ + model = _routed_oriented(accumulation) + mutate(model) + images = _cuda_flagged_images() + expected = model(images) + calls: list[None] = [] + model.backbone.register_forward_hook(lambda *_: calls.append(None)) + with ( + patch.object(model._network_capture, "run_oriented", wraps=model._network_capture.run_oriented) as spy, + pytest.warns(UserWarning, match=f"stays eager: {reason}"), + ): + output = model._training_network(images) + assert spy.call_count == recorder_calls + assert model._network_capture._disabled and model._network_capture._network is None + assert len(calls) == 1 + _same( + [getattr(expected, field.name) for field in fields(DualHeadOutput)], + [getattr(output, field.name) for field in fields(DualHeadOutput)], + ) + + +class _ClassificationOnly(DetectionLitModule): + """A supported native customization whose objective leaves angle gradients absent.""" + + def _task_extra_loss(self, context: Any, out: Any, stage: str) -> torch.Tensor: + """Keep classification connected while declining oriented supervision.""" + return torch.zeros_like(out.total) + + +def test_oriented_subclass_override_stays_native() -> None: + """A consumer overridden in a subclass retires replay, as an instance attribute does. + + The guard unwraps bound methods through ``__func__`` before comparing them. A subclass + method is a bound method whose function is not the native one, the path the + instance-level lambdas above never reach; this case runs it without a GPU. + """ + model = _ClassificationOnly(0.34, 0.25, 64, 2, task="obb", capture_network=True) + assert _oriented_replay_breach(model) == "the native oriented objective is overridden" + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA training") +def test_oriented_override_keeps_absent_angle_gradients(graph_recordings: MagicMock) -> None: + """The public trainer fallback skips angle state instead of materializing zero gradients. + + The replaying arm retires once, and for the overridden objective specifically: a + retirement for any other cause would leave the same absent gradients behind it. + """ + batches = _batches("uniform") + initial = _oriented().state_dict() + results = [] + for capture in (False, True): + model = _ClassificationOnly( + 0.34, + 0.25, + 64, + 2, + task="obb", + capture_network=capture, + w_muon=0.0, + w_sgd=1.0, + lr=1e-4, + warmup_epochs=0.0, + lrf=1.0, + ) + model.load_state_dict(initial) + trainer = Trainer( + accelerator="cuda", + devices=1, + max_epochs=1, + logger=False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + num_sanity_val_steps=0, + limit_val_batches=0, + gradient_clip_val=1.0, + deterministic=True, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + trainer.fit(model, train_dataloaders=DataLoader(copy.deepcopy(batches), batch_size=None)) + retired = [str(item.message) for item in caught if str(item.message).startswith("capture_network=True stays")] + assert len(retired) == int(capture) + assert all("the native oriented objective is overridden" in message for message in retired), retired + assert not graph_recordings.call_count + angles = [parameter for name, parameter in model.named_parameters() if "angle_stems" in name] + assert angles and all(parameter.grad is None for parameter in angles) + assert all(parameter not in trainer.optimizers[0].state for parameter in angles) + results.append(copy.deepcopy((model.state_dict(), trainer.optimizers[0].state_dict()))) + _same(*results)