diff --git a/docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md b/docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md new file mode 100644 index 00000000..20c67af8 --- /dev/null +++ b/docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md @@ -0,0 +1,453 @@ +# Extend compile_model_if_enabled to Radar, Image, and Audio Classification Implementation Plan + +**Goal:** Wire `compile_model_if_enabled` (the `torch.compile` warmup-and-fallback helper added by PR #22, currently only called from the four `timeseries_*` reference scripts) into `radar_classification`, `image_classification`, and `audio_classification`'s `main()` functions, so all seven reference training scripts get the same hardware-acceleration path — and measure whether it actually helps each one. + +**Architecture:** Each of the three target `main()` functions already has the exact structural shape `compile_model_if_enabled` expects: `move_model_to_device(model, device, logger)` immediately followed by `setup_distributed_model(model, args, device)`, with `dataset.X` (post feature-extraction) available in scope. The fix is the same one-line insertion in all three files, matching `timeseries_classification/train.py:275` verbatim: +```python +model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) +``` +placed between those two calls, plus adding `compile_model_if_enabled` to each file's import from `..common.train_base`. No changes to `compile_model_if_enabled` itself or to `train_base.py`. + +**Tech Stack:** Python 3.10, PyTorch 2.7.1, pytest + +## Global Constraints + +- Python `==3.10.*` +- No new dependencies +- Zero change to behavior when `--compile-model` is not set (the default, `0`) — `compile_model_if_enabled` already no-ops in that case; do not add any new conditionals around the call +- Do not touch `compile_model_if_enabled` itself, `train_base.py`, or any file outside the three target `train.py` files (+ their new/modified test files) +- Each task's benchmark must reuse the same measurement methodology as the radar entrypoint fix plan (`docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md`) — synthetic fixture, same device/epoch/batch-size reporting style — so results are comparable across all four now-benchmarked modules (timeseries already had compile; radar was benchmarked without it in that plan; this plan adds image and audio, and re-benchmarks radar with it now on) +- **Verification runs on two machines, not one:** a local macOS machine (CPU vs MPS) AND a separate CUDA-equipped Linux machine, referred to below as "GX10" (CPU vs CUDA — `NVIDIA GB10`, torch 2.9.0+cu130), reachable only by the controller over SSH. Implementer subagents are not expected to have SSH access to GX10 — the controller runs the GX10 leg directly after each task's local implementation is approved, and appends those results alongside the local ones in that task's Results section before considering the task's verification complete. + +## Context: Why This Is Needed + +While closing out the radar entrypoint fix (`docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md`), the whole-plan review found that fixing radar's `run()` dispatch (making `main()` live) did **not** close its MPS-vs-CPU performance gap, because `compile_model_if_enabled`/`apply_hardware_defaults` were never wired into radar's `main()` at all — only into the four `timeseries_*` scripts (PR #22, commit `dc8eeb2`). The same review flagged that `image_classification` and `audio_classification` have the identical gap: + +``` +grep -n "compile_model_if_enabled" tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py +tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py +``` +— zero matches in either file, confirmed directly against source, independent of the radar plan's own investigation. + +Rather than fix radar alone and leave image/audio in the same state, this plan closes the gap everywhere it exists in one coordinated pass, so the three currently-uncompiled reference scripts converge on the same behavior as `timeseries_classification`, `timeseries_regression`, `timeseries_forecasting`, and `timeseries_anomalydetection` already have. + +**Known unresolved question, not assumed:** whether `torch.compile` actually helps on hardware this small (radar's `LINEAR_4L_PC` is a 4-layer linear/BatchNorm model — the radar plan measured MPS at ~1.9x slower than CPU with no compile, and the theory is that fused kernels reduce per-op dispatch overhead, but this has not been confirmed for a model this size). Each task's benchmark step exists specifically to answer this per-module — report what you measure, not what the theory predicts. + +**Scope note (added after whole-plan review):** `compile_model_if_enabled` returns early, unmodified, whenever `args.quantization` is set (`train_base.py:727-739`) — all benchmarks in this plan used `--quantization 0` (the default), so every number below speaks only to float training; the QAT path is untouched by this wiring, by design. Separately, `PYTORCH_ENABLE_MPS_FALLBACK=1` (needed for `compile_model_if_enabled`'s MPS fallback warnings to be warnings rather than hard `NotImplementedError`s) is set via `os.environ.setdefault` at `tinyml-modelmaker`'s entry points, not inside `tinyml-tinyverse` itself — the supported path (through modelmaker) is unaffected, but a user invoking any of these `references/*/train.py` scripts directly with `--device mps --compile-model 1` will now hard-fail on an unsupported op instead of falling back, where the flag was previously inert (compile was never reachable at all). + +--- + +## File Map + +| Action | Path | Responsibility | +|--------|------|-----------------| +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` | Wire `compile_model_if_enabled` into `main()` | +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py` | Wire `compile_model_if_enabled` into `main()` | +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py` | Wire `compile_model_if_enabled` into `main()` | +| Create | `tinyml-tinyverse/tests/test_radar_compile_wired.py` | Regression test: radar's `main()` calls `compile_model_if_enabled` | +| Create | `tinyml-tinyverse/tests/test_image_classification_compile_wired.py` | Same, image classification | +| Create | `tinyml-tinyverse/tests/test_audio_classification_compile_wired.py` | Same, audio classification | + +--- + +## Task 1: Wire `compile_model_if_enabled` into radar_classification, with test + benchmark + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` (imports ~line 65-89; insertion point currently lines 251-254, between `move_model_to_device` and `setup_distributed_model`) +- Create: `tinyml-tinyverse/tests/test_radar_compile_wired.py` + +**Interfaces:** +- Consumes: `tinyml_tinyverse.references.common.train_base.compile_model_if_enabled` (existing, unmodified — signature `compile_model_if_enabled(model, args, logger, input_shape=None)`, returns the (possibly compiled) model) +- Produces: nothing new consumed by later tasks — Tasks 2 and 3 are independent, same pattern applied to different files + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: radar_classification.train.main() must call +compile_model_if_enabled, matching the pattern already used in the 4 +timeseries_* reference scripts (PR #22). Without this, --compile-model is +silently a no-op for radar regardless of what the caller requests.""" +import inspect + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +def test_main_calls_compile_model_if_enabled(): + main_source = inspect.getsource(radar_train.main) + assert "compile_model_if_enabled(" in main_source, ( + "main() does not call compile_model_if_enabled -- torch.compile/AMP " + "hardware acceleration would silently not apply to radar training." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_compile_wired.py -v` +Expected: FAIL — `"compile_model_if_enabled("` not present in `main`'s source + +- [x] **Step 3: Add the import and the call** + +In `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py`, add `compile_model_if_enabled` to the existing import from `..common.train_base` (around line 65-89 — find the multi-line `from ..common.train_base import (...)` block and add it as a new entry, alphabetically or grouped with `move_model_to_device`/`compile_model_if_enabled`-adjacent imports per the file's existing style). + +Then in `main()`, change: +```python + move_model_to_device(model, device, logger) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` +to: +```python + move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` + +- [x] **Step 4: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_compile_wired.py -v` +Expected: PASS + +- [x] **Step 5: Manual end-to-end sanity check with --compile-model 1** + +Using the synthetic radar fixture from the entrypoint-fix plan (regenerate via `make_radar_fixture.py` if not present — see that plan's Task 1 for the recipe), drive `train.run(args)` directly with `--device cpu --compile-model 1`, a few epochs. Confirm it completes without error and the log shows `compile_model_if_enabled`'s own INFO lines (check `train_base.py`'s `compile_model_if_enabled` for the exact log message text first — grep the run log for it). Then confirm `--compile-model 0` (the default) still behaves identically to before this change — no compile-related log lines, same training behavior. + +- [x] **Step 6: Benchmark CPU vs MPS with compile now enabled** + +Reuse the benchmark driver pattern from the entrypoint-fix plan (`bench_radar.py` in the session scratchpad), but add `--compile-model 1` to the argv. Run once on `cpu`, once on `mps` (`PYTORCH_ENABLE_MPS_FALLBACK=1`), same 30 epochs / batch size 16 / `LINEAR_4L_PC` as the entrypoint-fix plan's benchmark, so the numbers are directly comparable. Record: does compile change the CPU number, the MPS number, or the ratio between them, versus the entrypoint-fix plan's post-fix-no-compile numbers (CPU 0.621s/epoch, MPS 1.066s/epoch, 1.72x)? Report what you measure even if compile makes things worse or has no effect — that's a legitimate finding for a model this small, not a task failure. + +- [x] **Step 7: Record results and commit** + +Append a `## Results (Task 1: radar)` section to this plan doc with the benchmark table and a short explanation. + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/radar_classification/train.py tests/test_radar_compile_wired.py +git add ../docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md +git commit -m "feat: wire compile_model_if_enabled into radar_classification main()" +``` + +--- + +## Task 2: Wire `compile_model_if_enabled` into image_classification, with test + benchmark + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py` (imports ~line 92-111; insertion point currently lines 329-332, between `move_model_to_device` and `setup_distributed_model`) +- Create: `tinyml-tinyverse/tests/test_image_classification_compile_wired.py` + +**Interfaces:** +- Consumes: same `compile_model_if_enabled` as Task 1 — independent of Task 1, do not wait for it or reuse its branch + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: image_classification.train.main() must call +compile_model_if_enabled, matching the timeseries_* pattern (PR #22).""" +import inspect + +from tinyml_tinyverse.references.image_classification import train as image_train + + +def test_main_calls_compile_model_if_enabled(): + main_source = inspect.getsource(image_train.main) + assert "compile_model_if_enabled(" in main_source, ( + "main() does not call compile_model_if_enabled -- torch.compile/AMP " + "hardware acceleration would silently not apply to image classification training." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_image_classification_compile_wired.py -v` +Expected: FAIL + +- [x] **Step 3: Add the import and the call** + +Add `compile_model_if_enabled` to the existing `from ..common.train_base import (...)` block (~line 92-111). Then change: +```python + move_model_to_device(model, device, logger) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` +to: +```python + move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` +**Note the indentation in this file is one level deeper than radar/audio (this block sits inside an outer block in image_classification) — match the surrounding indentation exactly, don't copy radar's column position verbatim.** + +**Watch for:** this file has an `args.nn_for_feature_extraction` branch elsewhere (used at export time, line ~450-454) that chooses between `dataset.X_raw.shape` and `dataset.X.shape` depending on whether an NN is used for feature extraction. Confirm which shape the model actually consumes *at the point where you're inserting the compile call* (before or after any feature-extraction wrapping) — if `nn_for_feature_extraction` changes what the raw model's forward() expects at this point in `main()`, use the matching shape instead of assuming `dataset.X.shape` unconditionally. If unsure after reading the surrounding ~50 lines, ask before proceeding — this is exactly the kind of task-specific judgment call the brief can't resolve for you. + +- [x] **Step 4: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_image_classification_compile_wired.py -v` +Expected: PASS + +- [x] **Step 5: Manual end-to-end sanity check with --compile-model 1** + +Build or reuse a small synthetic image classification fixture (check `tinyml-tinyverse/tests/` for an existing image dataset test fixture/helper before building one from scratch — this repo likely already has one given image_classification has existing tests). Drive `train.run(args)` with `--device cpu --compile-model 1`, a few epochs. Confirm completion and compile-related log lines present. Confirm `--compile-model 0` still behaves as before. + +- [x] **Step 6: Benchmark CPU vs MPS with compile enabled** + +Same methodology as Task 1 Step 6, adapted to whichever image model this benchmark uses (pick the smallest/fastest registered image model available, to keep iteration time reasonable) and image_classification's own CLI args. This is the first CPU-vs-MPS benchmark for image classification in either plan — there's no prior "no-compile" baseline to compare against from the earlier work, so also run once with `--compile-model 0` on both devices first to get that baseline, then `--compile-model 1` to see the delta, in the same benchmark session. + +- [x] **Step 7: Record results and commit** + +Append `## Results (Task 2: image_classification)` to this plan doc. + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/image_classification/train.py tests/test_image_classification_compile_wired.py +git add ../docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md +git commit -m "feat: wire compile_model_if_enabled into image_classification main()" +``` + +--- + +## Task 3: Wire `compile_model_if_enabled` into audio_classification, with test + benchmark + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py` (imports ~line 92-111; insertion point currently lines 311-313, between `move_model_to_device` and `setup_distributed_model`) +- Create: `tinyml-tinyverse/tests/test_audio_classification_compile_wired.py` + +**Interfaces:** +- Consumes: same `compile_model_if_enabled` — independent of Tasks 1 and 2 + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: audio_classification.train.main() must call +compile_model_if_enabled, matching the timeseries_* pattern (PR #22).""" +import inspect + +from tinyml_tinyverse.references.audio_classification import train as audio_train + + +def test_main_calls_compile_model_if_enabled(): + main_source = inspect.getsource(audio_train.main) + assert "compile_model_if_enabled(" in main_source, ( + "main() does not call compile_model_if_enabled -- torch.compile/AMP " + "hardware acceleration would silently not apply to audio classification training." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_audio_classification_compile_wired.py -v` +Expected: FAIL + +- [x] **Step 3: Add the import and the call** + +Add `compile_model_if_enabled` to the existing `from ..common.train_base import (...)` block (~line 92-111). Then change: +```python + move_model_to_device(model, device, logger) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` +to: +```python + move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) +``` + +Same caveat as Task 2 applies here — audio_classification also has an `nn_for_feature_extraction` / `X_raw` vs `X` distinction at export time (line ~424-427). Check whether it affects what shape the model expects at this earlier insertion point before assuming `dataset.X.shape` unconditionally. + +- [x] **Step 4: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_audio_classification_compile_wired.py -v` +Expected: PASS + +- [x] **Step 5: Manual end-to-end sanity check with --compile-model 1** + +Same approach as Task 2 Step 5, adapted to audio_classification's dataset/CLI. Check `tinyml-tinyverse/tests/` for existing audio fixture helpers first. + +- [x] **Step 6: Benchmark CPU vs MPS with compile enabled** + +Same methodology as Task 2 Step 6 (baseline `--compile-model 0` then `--compile-model 1`, both devices), adapted to audio_classification's smallest registered model. + +- [x] **Step 7: Record results and commit** + +Append `## Results (Task 3: audio_classification)` to this plan doc. + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/audio_classification/train.py tests/test_audio_classification_compile_wired.py +git add ../docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md +git commit -m "feat: wire compile_model_if_enabled into audio_classification main()" +``` + +--- + +## Results (Task 1: radar) + +**Wiring:** `compile_model_if_enabled` is now imported and called in `radar_classification/train.py`'s `main()`, between `move_model_to_device` and `setup_distributed_model`, exactly matching the `timeseries_*` pattern. Regression test `tests/test_radar_compile_wired.py` passes; `--compile-model 0` (default) produces zero compile-related log lines and identical behavior to before this change; `--compile-model 1` produces `compile_model_if_enabled`'s own log line (`Compiling model with torch.compile (backend=aot_eager)`) on both CPU and MPS, with no fallback warnings on either device — compilation and the warmup forward pass succeed cleanly for `LINEAR_4L_PC` on this hardware. + +**Benchmark** — same fixture, `LINEAR_4L_PC`, batch size 16, 30 epochs, synthetic radar fixture, as the entrypoint-fix plan's post-fix-no-compile baseline: + +| Config | Device | s/epoch | CPU/MPS ratio | +|---|---|---|---| +| No compile (entrypoint-fix plan baseline) | CPU | 0.621 | 1.72x | +| No compile (entrypoint-fix plan baseline) | MPS | 1.066 | | +| `--compile-model 1` (this task) | CPU | 0.765 | 1.70x | +| `--compile-model 1` (this task) | MPS | 1.303 | | + +**Finding: `torch.compile` makes radar training slower on both devices, and does not close the CPU/MPS gap.** With compile enabled, CPU slows by ~23% (0.621s -> 0.765s/epoch) and MPS slows by ~22% (1.066s -> 1.303s/epoch). The CPU/MPS ratio is essentially unchanged (1.72x -> 1.70x) since both devices use the same `aot_eager` backend (the backend-selection logic in `compile_model_if_enabled` only routes CUDA to `inductor`; both CPU and MPS get `aot_eager`) and both slow down by a similar proportion. + +This is consistent with `LINEAR_4L_PC` being a tiny 4-layer linear/BatchNorm model: `torch.compile`'s per-call dynamo tracing, guard-checking, and graph-dispatch overhead is fixed cost per training step, and for a model this small there isn't enough per-op dispatch overhead in the eager path for kernel fusion to recoup that cost. The theory in this plan's "Known unresolved question" (fused kernels reducing dispatch overhead) does not hold for this model size — compile is a net loss here, not neutral or beneficial. Wiring it in is still correct (it makes `--compile-model` functional, matching the other reference scripts, and lets a caller opt in for larger/more compute-bound models where the tradeoff may differ), but the flag should not be turned on by default for radar's small linear models based on this evidence. + +**GX10 leg (controller-run, per Global Constraints):** same fixture/model/batch-size/epochs, GX10's Python venv (torch 2.9.0+cu130), `NVIDIA GB10`. GX10's clone was stale (`5035057`, pre-dating this whole plan) — updated to `origin/integration` (`e472749`) first. Environment needed several missing leaf packages installed (`--no-deps`, no `torch`/`torchvision` touched): `tabulate`, `torcheval`, `torchinfo`, `colorama`, `onnx`, `onnxruntime`, `protobuf`, `ml_dtypes`, `cryptography`, `PyWavelets`, `opencv-python`, `onnxscript`, `onnx_ir`. `cmsisdsp` and `torchaudio` were stubbed via `sys.modules` in the benchmark driver instead of installed for real — both are only needed transitively (by `timeseries_dataset.py`'s FFT path and `audio_dataset.py` respectively, eagerly imported by `datasets/__init__.py`) and never touched by radar's own code path; `cmsisdsp` in particular requires a multi-minute native CMSIS-DSP C build with no prebuilt ARM64 wheel, unrelated to what this benchmark measures. + +| Config | Device | s/epoch | +|---|---|---| +| No compile | CPU | 4.093 | +| No compile | CUDA | 0.856 | +| `--compile-model 1` | CPU | 3.972 | +| `--compile-model 1` | CUDA | 0.943 | + +CPU and CUDA aren't directly comparable to the Mac's CPU/MPS numbers (different, much less GX10-tuned CPU vs Apple Silicon CPU) — the useful comparison is each device against itself, with vs. without compile, on the same machine. + +**Finding: on GX10, CUDA beats CPU by 4.8x even with no compile** (0.856s vs 4.093s/epoch) — a completely different picture from the Mac, where MPS lost to CPU. GX10's GPU has enough throughput advantage over its CPU that it wins decisively on this tiny model without any fusion help. + +**Finding: `torch.compile` on CUDA never actually ran — it failed and silently fell back to eager.** The run log (`~/bench_radar/radar_out_cuda_compile1/run.log`) shows: +``` +INFO: root.main: Compiling model with torch.compile (backend=inductor) +WARNING: root.main: torch.compile failed (or failed its warmup pass), falling back to eager mode: +CalledProcessError: Command '['/usr/bin/gcc', '.../cuda_utils.c', '-O3', '-shared', '-fPIC', ... +-lcuda', '-L.../triton/backends/nvidia/lib', ...]' returned non-zero exit status 1. +``` +Triton's CUDA-kernel codegen fails to build `cuda_utils.c` via `gcc` on this box — a toolchain issue local to this environment (missing header/lib path for Triton's nvidia backend), not a `compile_model_if_enabled` defect. The `compile_model_if_enabled` warmup-and-fallback mechanism (built earlier this session, `docs/superpowers/plans/2026-07-28-compile-warmup-fallback.md`) caught the failure exactly as designed and fell back cleanly — training was not interrupted. This means the measured "compile=1, CUDA, 0.943s/epoch" number is **eager mode plus the one-time cost of a failed compile attempt**, not a real compiled-vs-uncompiled comparison. (**Update 2026-08-14:** the toolchain issue was fixed — see the Addendum in Task 3's Results section below for the real `inductor`-on-CUDA numbers.) + +**Finding: CPU's `aot_eager` backend did engage successfully on GX10** (`INFO: root.main: Compiling model with torch.compile (backend=aot_eager)`, no fallback warning) and produced a small apparent win (4.093 -> 3.972s/epoch, ~3%) — versus the Mac's own CPU result, a clear ~23% *loss* for the same model/backend. This is a single-shot measurement on both sides (see the whole-plan methodology caveat), and ~3% is well within the ~7-9% run-to-run wobble this session's benchmarking has otherwise documented, so treat this as "the two machines may not agree on aot_eager's sign for this model" rather than a confirmed win — it would need repeated runs on both machines to distinguish a real effect from noise. + +**Net implication for the "should compile be on by default" question:** no clean "yes" or "no" emerges. Mac CPU/MPS: compile hurts on both. GX10 CPU: a ~3% apparent win, single-shot and within this session's documented run-to-run noise — not confirmed real (see caveat below). GX10 CUDA: unmeasured at the time this was written — see the Addendum in Task 3's Results section for the real `inductor` numbers, obtained after the GX10 toolchain issue was fixed. Recommend leaving `--compile-model` opt-in (its current default); nothing in the full dataset (including the later `inductor` numbers) argues for changing it. + +--- + +## Results (Task 2: image_classification) + +**Wiring:** `compile_model_if_enabled` is now imported and called in `image_classification/train.py`'s `main()`, between `move_model_to_device` and `setup_distributed_model`, at the same indentation level as the surrounding `try:` block (one level deeper than radar/audio, per the brief's note). Regression test `tests/test_image_classification_compile_wired.py` passes; `--compile-model 0` (default) produces zero compile-related log lines and identical behavior to before this change; `--compile-model 1` produces `compile_model_if_enabled`'s own log line (`Compiling model with torch.compile (backend=aot_eager)`) on both CPU and MPS, with no fallback warnings on either device. + +**X vs X_raw judgment call:** at the insertion point (right after `move_model_to_device`), `model` is still the plain model returned by `models.get_model(args.model, variables, num_classes, input_features=input_features, ...)` (or `torch.load` for `--load-saved-model`), constructed from `variables = dataset.X.shape[1]` / `input_features = dataset.X.shape[2]` a few lines earlier (line 291-292) — confirmed by the existing `torchinfo.summary(model, summary_input_shape)` call at line 320-322, which also uses `dataset.X.shape[1:]` for this same pre-wrap model. The `nn_for_feature_extraction` / `X_raw` vs `X` branch the brief flagged only appears later (lines 342-350), where `model` gets wrapped in `NeuralNetworkWithPreprocess` — that wrapping happens *after* our insertion point, not before it. So regardless of `args.nn_for_feature_extraction`, the raw model being compiled at this point in `main()` always consumes `dataset.X.shape[1:]`-shaped input; the brief's suggested `input_shape=(1,) + dataset.X.shape[1:]` is correct unconditionally, and no branching on `nn_for_feature_extraction` was needed. + +**Fixture:** no ready-made image dataset fixture/driver existed in `tinyml-tinyverse/tests/` (the closest prior art, `tests/test_train_best_epoch_bugs_vision_audio.py`, fully mocks `get_model`/`create_data_loaders` and never drives real image data end-to-end). Built a small synthetic fixture instead: N class folders of 28x28 grayscale PNGs (`//*.png`, matching `GenericImageDataset`'s fallback-discovery layout — no `annotations/*_list.txt` needed since `--dataset` stays at its `folder` default, not `modelmaker`), each class centered on a distinct mean pixel intensity plus noise. Model: `CNN_LENET5`, the smallest of the three registered image models (~a few hundred params vs. `CNN_IMG_MOBILENETV1/2_58K_NPU`'s ~58K), whose spec hard-codes a 28x28x1 input (`Linear(in_features=400, ...)` after two conv+pool blocks) — fixture images were generated at exactly that size. Driver script (`bench_image.py` in the session scratchpad) builds argv for `train.get_args_parser().parse_args(argv)` and times `train.run(args)`, same pattern as `bench_radar.py`. + +**Manual E2E check (Step 5):** 3-class fixture (12 images/class), 2 epochs, batch size 8. `--device cpu --compile-model 0` (both pre- and post-code-change): completes cleanly, exports `model.onnx`, zero compile-related log lines. `--device cpu --compile-model 1`: completes cleanly, exports `model.onnx`, log shows `INFO: root.main: Compiling model with torch.compile (backend=aot_eager)` with no fallback warning. + +**Benchmark (Step 6)** — 5-class fixture (20 images/class = 100 images), `CNN_LENET5`, batch size 16, 30 epochs, `--device cpu` / `--device mps` (`PYTORCH_ENABLE_MPS_FALLBACK=1`), each config run 3-4 times to check stability: + +| Config | Device | s/epoch (steady-state avg) | compile engaged? | +|---|---|---|---| +| `--compile-model 0` (baseline) | CPU | 0.037 (n=3, range 0.037-0.037) | n/a (no compile lines in log) | +| `--compile-model 0` (baseline) | MPS | 0.073 (n=3 steady-state, range 0.069-0.077) | n/a (no compile lines in log) | +| `--compile-model 1` | CPU | 0.083 (n=3, range 0.083-0.084) | yes, `aot_eager`, no fallback | +| `--compile-model 1` | MPS | 0.094 (n=3 steady-state, range 0.092-0.096) | yes, `aot_eager`, no fallback | + +**Finding: `torch.compile` makes image_classification training slower on both devices, same direction as radar, but with a smaller relative hit on MPS than on CPU.** CPU regresses ~124% (0.037s -> 0.083s/epoch) — proportionally worse than radar's ~23% CPU regression. MPS regresses only ~29% (0.073s -> 0.094s/epoch) — close in magnitude to radar's ~22% MPS regression, but here the CPU/MPS gap actually **narrows** with compile on: 0.073/0.037 = 1.98x (MPS slower, no compile) vs. 0.094/0.083 = 1.13x (MPS slower, with compile). CPU still wins in absolute terms either way, but compile shrinks rather than preserves the ratio — the opposite of radar, where the ratio stayed essentially flat (1.72x -> 1.70x). This is driven by the CPU/MPS *asymmetry* (CPU regresses far more than MPS does), not by MPS itself improving — no causal mechanism for that asymmetry was established here (see the whole-plan review's note that a competing "conv work narrows the GPU's relative advantage" story doesn't hold up against the GX10 data below, so this is reported as an observed pattern, not an explained one). + +**Correction (2026-08-14, per independent peer review): the numbers above conflate one-time compile warmup with steady-state cost, and the "net loss" conclusion below does not survive isolating them.** `s/epoch` here is whole-`train.run()` time ÷ 30, and Task 3's own MPS-dropout finding already showed compile's one-time warmup (~1.4s there) can be large relative to a 30-epoch total when each epoch is fast. Re-ran with the 1-epoch-vs-30-epoch isolation Task 3 used, fresh runs, same fixture/model/batch-size: + +| Config | Device | 1-epoch total | 30-epoch total | steady-state (solved) | fresh no-compile baseline | +|---|---|---|---|---|---| +| `--compile-model 1` | CPU | 1.519s | 1.860s | ~0.012s/epoch (warmup ~1.51s) | 0.019s/epoch | +| `--compile-model 1` | MPS | 1.310s | 2.011s | ~0.024s/epoch (warmup ~1.29s) | 0.048s/epoch | + +Once warmup is excluded, compiled steady-state is **faster than eager on both devices** — roughly 37% faster on CPU, roughly 50% faster on MPS — the opposite sign from the naive 30-epoch-averaged "CPU regresses ~124%, MPS regresses ~29%" finding above. The original numbers weren't wrong as measurements of a 30-epoch run; they're wrong as a proxy for steady-state training cost, because a 30-epoch run this fast (~0.6-1.4s total) barely amortizes a ~1.3-1.5s one-time compile cost. (These are fresh runs in a new session, not re-runs of the exact original data — but the no-compile baselines were measured fresh in the same pass for a fair within-session comparison.) This reverses the "Bottom line" below for any training run long enough to amortize the warmup — see the whole-plan synthesis for how this changes the cross-task picture. + +**Methodology note — MPS cold-start artifact:** the very first MPS invocation of the whole benchmark session (`--compile-model 0`, run 1) measured 0.171s/epoch, 2.3-4.6x slower than every subsequent MPS run (compile-0 or compile-1) in the same session, each a fresh process. All later MPS runs (4 for compile-0, 4 for compile-1, alternating) clustered tightly (0.069-0.077 and 0.090-0.096 respectively). This looks like a one-time cost paid on the first-ever MPS/Metal invocation in the session (e.g. Metal shader compilation cache being cold on disk, not per-process) rather than genuine compile-vs-no-compile signal, since it appeared on a `--compile-model 0` run. The table above reports the steady-state average (excluding that one outlier); the raw run-by-run numbers are preserved in the session scratchpad driver output for reference. This is the same class of measurement caveat Task 1 flagged for GX10 (environment-specific first-run cost, not part of the compile question itself) — worth keeping in mind for any future single-shot MPS benchmark on this or other modules. + +**Bottom line (revised 2026-08-14):** wiring `compile_model_if_enabled` into `image_classification` is correct and makes `--compile-model` functional for this script, matching the other reference scripts and Task 1's radar wiring. Unlike the original conclusion, the warmup-isolated evidence does *not* argue against compile for `CNN_LENET5` — steady-state throughput improves on both CPU and MPS. It does argue against `--compile-model` defaulting on for *short* runs (a handful of epochs), where the one-time warmup cost dominates and compile is a net loss in wall-clock terms even though it's faster per-epoch once warmup is paid. Whether `--compile-model` should default on depends entirely on typical training-run length, which this benchmark's fixture (a few dozen epochs on a tiny synthetic dataset) doesn't represent — real training runs are likely long enough to amortize the ~1.3-1.5s warmup easily. A larger, more compute-bound image model (e.g. `CNN_IMG_MOBILENETV1_58K_NPU`/`CNN_IMG_MOBILENETV2_58K_NPU`) was not benchmarked here. + +**GX10 leg (controller-run):** same fixture/model/batch-size, 30 epochs, GX10's Python venv (torch 2.9.0+cu130), `NVIDIA GB10`. GX10's clone fast-forwarded cleanly to `bf261a7` this time (no reset needed, unlike Task 1's stale-clone situation). Same `cmsisdsp`/`torchaudio` `sys.modules` stubs as Task 1 (still unrelated to image_classification's own code path); no new missing packages this time — everything Task 1 installed already covered it. + +| Config | Device | s/epoch | compile engaged? | +|---|---|---|---| +| No compile | CPU | 0.097 | n/a | +| No compile | CUDA | 0.074 | n/a | +| `--compile-model 1` | CPU | 0.198 | yes, `aot_eager`, no fallback | +| `--compile-model 1` | CUDA | 0.113 | attempted `inductor`, failed, fell back to eager (same Triton/gcc `cuda_utils.c` build failure as Task 1) | + +**Finding: CUDA beats CPU on GX10 even without compile, but only by 1.3x here** (0.074 vs 0.097s/epoch) — a much smaller margin than radar's 4.8x. Note this doesn't hold up as a simple "more conv work narrows the GPU's edge" story once Task 3's audio numbers are in: audio (also conv-heavy) shows a 2.2x GX10 GPU advantage, between radar's pure-linear 4.8x and image's 1.3x — the ordering is radar > audio > image, not monotonic in how conv-heavy the model is. Reported as an observed number, not an explained one. + +**Finding: `aot_eager` on CPU is a bigger loss on GX10 than on the Mac for this model** — 0.097 -> 0.198s/epoch is a ~104% regression on GX10's CPU, roughly consistent in direction and rough magnitude with the Mac's own CPU regression for image_classification (~124%) and both are much worse than either machine's CPU regression for radar. Small-CNN `aot_eager` compilation overhead appears to be a bigger relative cost than small-linear-net overhead was, on both machines. + +**Finding: the same CUDA/Triton/gcc build failure from Task 1 reproduces identically here** — confirms it's an environment-level GX10 toolchain issue (Triton's nvidia backend `cuda_utils.c` failing to compile via `gcc`), not something specific to radar's model or code path. (**Update 2026-08-14:** fixed — see the Addendum in Task 3's Results section.) + +--- + +## Results (Task 3: audio_classification) + +**Wiring:** `compile_model_if_enabled` is now imported and called in `audio_classification/train.py`'s `main()`, between `move_model_to_device` and `setup_distributed_model`, at the same indentation/position as radar's wiring (no extra `try:` nesting, unlike image_classification). Regression test `tests/test_audio_classification_compile_wired.py` passes; `--compile-model 0` (default) produces zero compile-related log lines and identical behavior to before this change; `--compile-model 1` produces `compile_model_if_enabled`'s own log line (`Compiling model with torch.compile (backend=aot_eager)`) on both CPU and MPS, with no fallback warnings on either device. + +**X vs X_raw judgment call:** same resolution as Task 2, re-verified by reading `main()` directly rather than assumed. At the insertion point (right after `move_model_to_device`), `model` is still the plain model returned by `models.get_model(args.model, variables, num_classes, input_features=input_features, ...)` (or `torch.load` for `--load-saved-model`), built from `variables = dataset.X.shape[1]` / `input_features = tuple(dataset.X.shape[2:])` a few lines earlier (line 274-276) — confirmed by the existing `torchinfo.summary(model, summary_input_shape)` call at line 302-304 (gated on `args.generic_model or args.nas_enabled`), which also uses `dataset.X.shape[1:]` for this same pre-wrap model. The `nn_for_feature_extraction` / `X_raw` vs `X` branch the brief flagged only appears later (lines 322-331), where `model` gets wrapped in `NeuralNetworkWithPreprocess` (and, when `args.nn_for_feature_extraction` is set, a separately-trained `FEModelLinear` feature extractor consuming `dataset.X_raw`) — that wrapping happens *after* our insertion point, not before it, exactly as in Task 2. The `X_raw`-shaped input only matters again much later, at export time (line 424, `input_shape = (1,) + dataset.X_raw.shape[1:]` when `nn_for_feature_extraction` is set) — a separate, later insertion point this task doesn't touch. So regardless of `args.nn_for_feature_extraction`, the raw model being compiled at this point in `main()` always consumes `dataset.X.shape[1:]`-shaped input; the brief's `input_shape=(1,) + dataset.X.shape[1:]` is correct unconditionally, no branching needed. + +**Fixture:** no existing audio fixture/driver in `tinyml-tinyverse/tests/` drives real audio data end-to-end (the closest prior art, `tests/test_train_best_epoch_bugs_vision_audio.py`, fully mocks `get_model`/`create_data_loaders`/etc. for both image and audio and never touches a real `GoogleSpeechCommandsDataset`). Built a small synthetic fixture instead: N class folders of 1-second, 16kHz mono WAV clips (`//*.wav`, matching `GoogleSpeechCommandsDataset`'s fallback-discovery layout — `glob(/*/*.wav)`, no `annotations/*_list.txt` needed since `--dataset` stays at its `folder` default), each class a distinct sine-tone frequency (200-3000 Hz, evenly spaced) plus Gaussian noise, generated with `soundfile` (confirmed `torchaudio.load` reads these back correctly with the installed `soundfile` backend — `torchaudio.list_audio_backends()` returns `['soundfile']` in this env). Model: `CNN_AUDIO_DSCNN`, the *only* model registered for `audio_classification` in `tinyml-modelzoo` (`tinyml_modelzoo/models/audio.py`) — a depthwise-separable-conv DSCNN (7 `Conv2d` layers, several depthwise, BatchNorm/ReLU/Dropout between them, ~23,171 parameters at default `filters=64`), a materially different op mix from radar's pure linear/BatchNorm stack and image's single small non-separable CNN. Its documented expected input is `(N, 1, 49, 10)`; the default audio CLI params (`sample_rate=16000`, `audio_duration_ms=1000` -> 16000 samples, `n_mfcc=10`, `frame_length_ms=30` -> `n_fft=480`, `frame_step_ms=20` -> `hop_length=320`) produce exactly 49 MFCC time frames x 10 coefficients with `center=False`, so no CLI overrides were needed to hit the model's expected shape. Driver script (`bench_audio.py` in the session scratchpad) builds argv for `train.get_args_parser().parse_args(argv)` and times `train.run(args)`, same pattern as `bench_radar.py`/`bench_image.py`. + +**Methodology finding — `--sampling-rate` vs `--sample-rate` naming collision (pre-existing, not part of this task's fix):** while building the fixture, an initial run with `--sampling-rate 1.0` (the value Tasks 1/2 used for the base parser's generic, audio-irrelevant, `required=True` FFT `--sampling-rate` arg) crashed inside `torchaudio`'s MFCC transform with `RuntimeError: stft(...) : expected 0 < n_fft < 1, but got n_fft=0`. Root cause: `GoogleSpeechCommandsDataset.__init__` does `for key, value in kwargs.items(): setattr(self, key, ...)` over every CLI arg passed via `**vars(args)`, then reads `self.sampling_rate = int(getattr(self, "sampling_rate", 16000))` — but the *base* parser's generic `--sampling-rate` (dest `sampling_rate`, meant for radar/timeseries FFT preprocessing) collides with and silently overwrites the dataset's own `sampling_rate` attribute, while `audio_classification`'s own `--sample-rate` flag (dest `sample_rate`, no "ing") is never read by the dataset at all — it sets an unused attribute. With `--sampling-rate 1.0`, the dataset computed `n_fft = int(1 * 30 / 1000) = 0`, crashing MFCC extraction. Worked around by passing `--sampling-rate 16000` (the real audio sample rate) instead of `1.0` for this benchmark. This is a genuine, pre-existing naming/wiring bug in the audio dataset loader unrelated to `compile_model_if_enabled` and outside this task's file list (`audio_dataset.py`, not `audio_classification/train.py`) — not fixed here, flagged for separate follow-up. + +**Manual E2E check (Step 5):** 3-class fixture (12 clips/class), 2 epochs, batch size 16, `--sampling-rate 16000`. `--device cpu --compile-model 0` (both pre- and post-code-change): completes cleanly, exports `model.onnx`, zero compile-related log lines. `--device cpu --compile-model 1`: completes cleanly, exports `model.onnx`, log shows `INFO: root.main: Compiling model with torch.compile (backend=aot_eager)` with no fallback warning. `--device mps --compile-model 0` and `--compile-model 1` (`PYTORCH_ENABLE_MPS_FALLBACK=1`): both complete cleanly and export `model.onnx`; compile-1 shows the same `aot_eager` compile line, no `compile_model_if_enabled` fallback warning, but does show a PyTorch `UserWarning` that `aten::native_dropout` isn't supported on MPS and falls back to CPU for that op — see the benchmark finding below. + +**Benchmark (Step 6)** — 5-class fixture (20 clips/class = 100 clips), `CNN_AUDIO_DSCNN`, batch size 16, 30 epochs, `--device cpu` / `--device mps` (`PYTORCH_ENABLE_MPS_FALLBACK=1`), each config run 3-4 times to check stability (whole-`train.run()` timing, same methodology as Tasks 1-2, so includes one-time dataset-load/MFCC-extraction and ONNX-export overhead amortized over 30 epochs): + +| Config | Device | s/epoch (avg) | compile engaged? | +|---|---|---|---| +| `--compile-model 0` (baseline) | CPU | 1.242 (n=3, range 1.222-1.254) | n/a (no compile lines in log) | +| `--compile-model 0` (baseline) | MPS | 0.110 (n=4, range 0.108-0.114) | n/a (no compile lines in log) | +| `--compile-model 1` | CPU | 1.186 (n=3, range 1.170-1.204) | yes, `aot_eager`, no fallback | +| `--compile-model 1` | MPS | 0.207 (n=4, range 0.201-0.211) | yes, `aot_eager`, no fallback (but see dropout finding below) | + +**Finding: audio_classification matches neither radar's nor image's pattern — MPS massively outperforms CPU here, the opposite of both prior tasks, and `torch.compile` is a small net win on CPU but a large net loss on MPS.** Without compile, MPS is ~11.3x *faster* than CPU (0.110s vs 1.242s/epoch) — a complete inversion of radar (MPS 1.7x slower) and image (MPS ~2x slower). A quick 1-epoch-vs-30-epoch differential run confirms this isn't a fixed-cost artifact: isolating the one-time dataset-load overhead (~0.28s CPU, ~0.58s MPS) from steady-state per-epoch cost still gives CPU ≈1.21s/epoch vs MPS ≈0.09s/epoch, a ~13x native throughput gap. `CNN_AUDIO_DSCNN` is a real, moderately-sized conv-heavy model (23,171 params, 7 `Conv2d` layers incl. depthwise) operating on a 49x10 feature map at batch 16, giving Apple's GPU real work to do. But "more conv work means the GPU wins more" doesn't hold up as a general story: GX10's own GPU-vs-CPU ratios for the identical three models are radar 4.8x > audio 2.2x > image 1.3x, not ordered by conv-heaviness. An equally plausible reading of this specific 11.3x gap is that Mac CPU is unusually slow at this workload rather than MPS being unusually fast at it — Mac CPU's audio number (1.242s/epoch) is 33.6x its own image number, while GX10 CPU's audio number is only 1.8x its own image number, a ~19x discrepancy in relative CPU behavior between the two machines for the same model. Both explanations are consistent with the data; this plan doesn't have enough evidence to pick one. + +With compile enabled, CPU improves slightly (1.242 -> 1.186s/epoch, ~-4.5%) — the only device across all three tasks where `torch.compile` was a net *win*, however small — while MPS regresses sharply (0.110 -> 0.207s/epoch, ~+88%). This **narrows** the CPU/MPS gap the same direction as image_classification did (11.3x -> 5.7x) but starting from the opposite baseline (MPS ahead, not behind). The MPS regression has a concrete, observed mechanism rather than just inferred dynamo/guard overhead: the `--compile-model 1` MPS log (and only that log — the `--compile-model 0` MPS log has zero such warnings, despite the model using the same `Dropout(0.2)`/`Dropout(0.4)` layers in both cases) shows `UserWarning: The operator 'aten::native_dropout' is not currently supported on the MPS backend and will fall back to run on the CPU`. This means compiling routes dropout through a lower-level decomposition (`aten::native_dropout`, likely via dynamo/AOTAutograd's graph capture) that lacks an MPS kernel, forcing a real CPU round-trip on every forward pass in the compiled path — whereas eager mode's `F.dropout` apparently dispatches through a path with native MPS support. A finer-grained 1-epoch-vs-30-epoch differential on the compiled MPS runs suggests the one-time `torch.compile` warmup cost (~1.4s) is also non-trivial relative to the 30-epoch total, so the naive 30-epoch-average ~88% regression somewhat overstates the true steady-state per-epoch cost (~56% by that isolation) — still a clear, large loss either way, just with the same one-time-warmup-dilution caveat Task 2 flagged for its MPS cold-start artifact. + +**Bottom line:** wiring `compile_model_if_enabled` into `audio_classification` is correct and makes `--compile-model` functional for this script, matching all six other reference scripts. Unlike radar and image_classification, the evidence here does *not* uniformly argue against compile: CPU sees a small, consistent win, and the MPS regression has an identified, addressable-in-principle cause (a dropout decomposition lacking an MPS kernel under `torch.compile`) rather than being an unavoidable property of small-model dynamo overhead. Still, the MPS loss is large enough (and CPU's win small enough) that there's no case for flipping the opt-in default here either — `--compile-model` should stay opt-in, as it does today, but audio_classification is the first of the three tasks where "does compile help" doesn't have a uniformly negative answer, and the dropout/MPS-kernel angle is a concrete, non-speculative lead if `torch.compile` + MPS support for this workload is ever revisited. + +**GX10 leg (controller-run):** same fixture/model/batch-size, 30 epochs, GX10's Python venv (torch 2.9.0+cu130), `NVIDIA GB10`. Clone fast-forwarded cleanly to `faad516`. Environment needed real (not stubbed) `torchaudio` this time, since `GoogleSpeechCommandsDataset` genuinely calls `torchaudio.load`/`MFCC`/`resample` (unlike radar/image, where `cmsisdsp`/`torchaudio` were only pulled in as unused eager imports and safely stubbed via `sys.modules`). `pip install torchaudio` grabbed 2.11.0 by default — ABI-incompatible with torch 2.9.0 (`undefined symbol: torch_library_impl` on import). Pinning `torchaudio==2.9.0` matched cleanly, but its default audio backend (`torchcodec`) requires system FFmpeg, which isn't installed on GX10 and isn't worth adding just for this benchmark. Worked around by monkey-patching `torchaudio.load` in the GX10 driver script only (not the repo) to read via the already-installed `soundfile` instead. + +This patch does feed `_load_audio`'s MFCC feature extraction (`audio_dataset.py:280-305`), so it does affect `dataset.X`, not just I/O plumbing — the earlier claim that it "doesn't affect what's being measured" was imprecise. What actually protects the measurement: `_load_audio` unconditionally casts to `float32` (neutralizing `soundfile`'s float64 default) and peak-normalizes every clip to max-amplitude 1.0 (neutralizing any amplitude-scale difference from `torchaudio.load`'s own normalization), and since the fixture and the `--sampling-rate` arg both use 16kHz, the resampling branch never fires on either loader. The one risk that doesn't self-correct is channel orientation (`soundfile` returns frames-first, `torchaudio.load` returns channels-first) — a wrong-axis patch would silently collapse the mono-downmix step to a near-zero waveform with no exception. Checked empirically rather than left as an argument: the GX10 audio run's `run.log` shows `Acc@1 100.000` at the best epoch — well above the ~20% chance level for this 5-class fixture, confirming the model learned real, distinguishable signal and the patch was not silently corrupting input. (Separately: `s/epoch` timing, which is what's actually reported in the table below, is driven by tensor shapes fixed by the MFCC parameters, not by waveform content — so even a corrupted patch would not have moved the reported numbers, though it could have invalidated the accuracy figure.) + +| Config | Device | s/epoch | compile engaged? | +|---|---|---|---| +| No compile | CPU | 0.174 | n/a | +| No compile | CUDA | 0.079 | n/a | +| `--compile-model 1` | CPU | 0.305 | yes, `aot_eager`, no fallback | +| `--compile-model 1` | CUDA | 0.141 | attempted `inductor`, failed (same Triton/gcc `cuda_utils.c` build failure as Tasks 1 and 2), fell back to eager | + +**Finding: CUDA beats CPU 2.2x on GX10 without compile** (0.079 vs 0.174s/epoch) — between radar's 4.8x and image's 1.3x, consistent with `CNN_AUDIO_DSCNN`'s real conv workload giving the GPU a clear but not overwhelming edge over GX10's capable CPU. + +**Finding: unlike the Mac, GX10's CPU `aot_eager` backend is a clear net LOSS for audio (+75%, 0.174 -> 0.305s/epoch) — the opposite sign from the Mac's small CPU win (~-4.5%).** Radar showed the same kind of Mac/GX10 CPU sign disagreement (Mac +23% loss vs GX10 ~3% apparent win — see that Results section's caveat that the GX10 number is single-shot and within noise). Audio's flip is the more solid of the two: -4.5% and +75% are both large enough, relative to this session's ~7-9% documented run-to-run wobble, to treat as more than noise on both sides, even without repeated runs. Since GX10's CUDA compile attempt never actually ran (Triton/gcc failure, same as every other module on this box), there's no GX10 MPS-dropout-style mechanism to check on the GPU side — CUDA does have a native dropout kernel, so the MPS-specific finding from the Mac leg doesn't apply here regardless. + +**Finding: the Triton/gcc `inductor` build failure reproduces identically for the third module in a row** — now confirmed systemic across all of radar, image_classification, and audio_classification on this GX10 environment, not tied to any particular model or code path. + +**Revised 2026-08-14, following independent peer review and the image warmup-isolation re-check above.** The original version of this section claimed "`aot_eager` is a consistent net loss... across every device class actually measured, with two exceptions near the noise floor." That claim doesn't survive scrutiny — one of the "exceptions" was itself mischaracterized, and two of the seven "regressions" turned out to be warmup-dilution artifacts, not steady-state cost, once actually isolated. + +**All nine naive 30-epoch-averaged `aot_eager` deltas** (the raw numbers as originally measured, before any warmup isolation): + +| module | Mac CPU | Mac MPS | GX10 CPU | +|---|---|---|---| +| radar | +23% | +22% | -3.0% (single-shot, within documented ~7-9% noise — not confirmed real) | +| image | +124% | +29% | +104% | +| audio | -4.5% (n=3, non-overlapping ranges — real) | +88% (~56% steady-state, isolated) | +75% | + +**What changes when warmup is actually isolated (only done for image and audio's MPS leg — radar and the other legs were not re-checked this way):** image's naive CPU/MPS deltas (+124%/+29%) reverse sign entirely once the ~1.3-1.5s one-time compile cost is excluded — steady-state compile is ~37% faster on CPU and ~50% faster on MPS (see the warmup-isolation correction in Task 2's Results above). Audio's Mac MPS delta shrinks from +88% (naive) to ~+56% (steady-state) but stays a large real loss. Audio's Mac CPU -4.5% was run n=3 with non-overlapping ranges (1.222-1.254s vs 1.170-1.204s) — a small but genuine, controlled result, not noise; treating it as "near the noise floor" (the original version of this section did) was itself an error, since a noise-floor figure derived from single uncontrolled runs elsewhere in this plan was applied to discount a result that was actually more rigorously measured than the figure used to discount it. + +**Net honest picture:** there is no clean "compile helps" or "compile hurts" story, and — unlike the original version of this section claimed — the data doesn't support a "consistent net loss" story either. What actually holds up: (1) radar and audio's Mac MPS regressions are large and were not explained away by warmup (audio's mechanism — a dropout op falling back to CPU under compile — is a real, identified cause); (2) image's apparent regressions were mostly or entirely warmup dilution and reverse to a real win in steady state; (3) whether `--compile-model` helps depends heavily on model architecture, device, and training-run length, none of which this benchmark holds constant enough to generalize from. `inductor`-on-CUDA is no longer unmeasured — see the Addendum immediately below, obtained after the GX10 toolchain issue was fixed. `--compile-model` staying opt-in (its current default) is still the right call given how mixed and run-length-dependent this data is — but not because "compile is a consistent loss." + +**Addendum (2026-08-14): the GX10 `inductor` block is resolved, all three modules re-measured.** Root cause: GX10's system Python 3.12 was missing its `python3.12-dev` package, so `Python.h` didn't exist anywhere on the box — Triton's `cuda_utils.c` needs it to build the CUDA extension it compiles per-kernel. Fixed by installing `python3.12-dev` (`sudo apt install python3.12-dev`, done by the repo owner directly on GX10 — outside what this session could do without their password). Re-ran the `--compile-model 1`/CUDA config for all three modules; `inductor` now builds and engages with zero fallback warnings on all three (previously it failed and silently fell back to eager every time, so none of the numbers below existed until now): + +| module | CUDA no compile | CUDA `inductor` | delta | +|---|---|---|---| +| radar | 0.856s/epoch | 1.062s/epoch | +24% | +| image | 0.074s/epoch | 0.375s/epoch | +407% | +| audio | 0.079s/epoch | 0.374s/epoch | +373% | + +`inductor` is a large regression for all three at 30 epochs — likely dominated by one-time compilation/kernel-autotuning cost rather than steady-state throughput: radar's no-compile run alone takes ~26s for 30 epochs (enough to amortize some fixed cost), while image's and audio's no-compile runs take only ~2.2-2.4s total for 30 epochs, so a multi-second one-time `inductor` autotuning pass would dominate their totals almost entirely — the same warmup-dilution effect Task 2 and Task 3 already documented for MPS cold-start and compile-warmup costs, here likely far larger in absolute terms since `inductor` autotunes real CUDA kernels rather than `aot_eager`'s lighter graph capture. This wasn't isolated with a longer run (out of scope for this addendum) — reported as an observed 30-epoch-total number with this caveat, not a steady-state claim. Directionally, this means `inductor`-on-CUDA does not obviously help these three small models any more than `aot_eager` did elsewhere in this plan, and for very short runs it can hurt substantially more. + +**Methodology caveat:** rigor was not uniform across this plan. Task 1's Mac CPU/MPS numbers are a single run each, compared against a single run from a separate prior plan and session — before Task 2 discovered that the first MPS invocation of a session can run 2.3-4.6x slower than steady state, and before Tasks 2/3 adopted an n=3-4-repeats-with-range methodology. Task 1's Mac numbers were never re-validated under that later, stricter methodology. All six GX10 data points (across all three tasks) are single-shot, with GX10's own run-to-run variance never characterized. This doesn't overturn any of the seven large-regression findings, but it's the reason the two sub-5% deltas (radar GX10 CPU, audio Mac CPU) are treated as "near the noise floor" above rather than as confirmed real effects. diff --git a/docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md b/docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md new file mode 100644 index 00000000..c8e01120 --- /dev/null +++ b/docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md @@ -0,0 +1,181 @@ +# Radar Training Entrypoint Fix Implementation Plan + +**Goal:** Wire the radar classification training script's `run()` entrypoint to the fully-featured `main()` function instead of the leftover `main_debug()` harness, then confirm on real hardware that this closes the MPS-vs-CPU performance gap it currently causes. + +**Architecture:** `run_distributed(main_debug, args)` becomes `run_distributed(main, args)` in `radar_classification/train.py`. `main()` already exists and is fully wired to the shared, hardened training infrastructure (`quantization_wrapped_model`, `compile_model_if_enabled`/`apply_hardware_defaults`, `resume_from_checkpoint`, `create_data_loaders`) — it is simply never called today. No new code is needed for the fix itself; the risk is regression in radar-specific behavior that only `main_debug` currently exercises, which Task 1's test and manual run guard against. + +**Tech Stack:** Python 3.10, PyTorch 2.7.1, pytest + +> **Correction (see ## Results below):** the Goal and Architecture above assumed `main()` was already wired to `compile_model_if_enabled`/`apply_hardware_defaults`. It is not. Task 1's fix is real (quantization and checkpoint-resume now work) but does **not** close the MPS-vs-CPU gap — see Results for the actual post-fix numbers and root cause. + +## Global Constraints + +- Python `==3.10.*` +- No new dependencies +- `main_debug()` itself is left in place (unreferenced) — it's the author's own notebook-parity debug harness per its docstring, not this plan's concern to delete +- The fix must not change CLI argument names or defaults — only which internal function `run()` dispatches to + +## Context: Why This Is Needed + +`BaseRadarModelTraining.run()` (`tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py`) calls `self.train_module.run(args)` for both the float-training pass and the quantized-training pass. `train_module` is wired to `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py`, whose `run()` is: + +```python +def run(args): + """Run training with optional distributed mode.""" + run_distributed(main_debug, args) +``` + +`main_debug()` is a hand-rolled duplicate of `main()` (its own comment: *"Following as close as possible steps from jupyter notebook to test if model learning plateau is coming from training loop"*). Confirmed by reading both functions side by side: + +- `main_debug()` hardcodes `phase = 'FloatTrain'` unconditionally and never calls `quantization_wrapped_model` — the "QuantTrain" argv pass `radar_base.py` builds produces a second float-training run instead, logged under the wrong phase label, so `get_radar_classification_log_summary_regex()`'s `QuantTrain`-specific regexes never match anything. +- `main_debug()` never calls `compile_model_if_enabled` or `apply_hardware_defaults` (PRs #21–23) — no AMP, no `torch.compile`, regardless of device or hardware. +- `main_debug()` never calls `resume_from_checkpoint`. + +Empirically confirmed this has a real, measurable performance cost, not just a theoretical one. Benchmarked `LINEAR_4L_PC` (the one registered radar model) on a synthetic 5-class radar-shaped fixture, 30 epochs, via the actual `radar_classification.train.run(args)` entrypoint as currently shipped: + +| Device | Time/epoch | +|---|---| +| CPU | 0.602s | +| MPS | 1.145s (**1.90x slower than CPU**) | + +MPS being slower than CPU is the expected signature of a small-op-heavy graph with no kernel fusion — exactly what `apply_hardware_defaults`/`compile_model_if_enabled` exists to fix elsewhere in the codebase. Task 2 re-runs this same benchmark once `main()` is live, to confirm and quantify the improvement. + +--- + +## File Map + +| Action | Path | Responsibility | +|--------|------|-----------------| +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py:522-524` | `run()` dispatches to `main`, not `main_debug` | +| Create | `tinyml-tinyverse/tests/test_radar_entrypoint_uses_main.py` | Regression test pinning the dispatch target | + +--- + +## Task 1: Fix `run()` dispatch target, with regression test + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` (function `run`, currently lines 522-524) +- Create: `tinyml-tinyverse/tests/test_radar_entrypoint_uses_main.py` + +**Interfaces:** +- Consumes: `tinyml_tinyverse.references.radar_classification.train.main`, `.main_debug`, `.run`, `.run_distributed` (all already defined in the module — no new interfaces) + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: radar_classification.train.run() must dispatch to main(), +not main_debug() (a leftover notebook-parity harness that silently skips +quantization, AMP, and torch.compile for every radar training run).""" +from unittest.mock import patch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +def test_run_dispatches_to_main_not_main_debug(): + with patch.object(radar_train, "run_distributed") as mock_run_distributed: + fake_args = object() + radar_train.run(fake_args) + + mock_run_distributed.assert_called_once() + dispatched_fn, dispatched_args = mock_run_distributed.call_args[0] + + assert dispatched_fn is radar_train.main, ( + f"run() dispatched to {dispatched_fn.__name__!r}, expected 'main'. " + "main_debug() never applies quantization_wrapped_model, " + "compile_model_if_enabled/apply_hardware_defaults, or " + "resume_from_checkpoint -- wiring run() to it silently drops quantized " + "training and hardware acceleration for every radar run." + ) + assert dispatched_args is fake_args +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_entrypoint_uses_main.py -v` +Expected: FAIL — `dispatched_fn` is `main_debug`, not `main` + +- [x] **Step 3: Fix the dispatch target** + +In `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py`, change: + +```python +def run(args): + """Run training with optional distributed mode.""" + run_distributed(main_debug, args) +``` + +to: + +```python +def run(args): + """Run training with optional distributed mode.""" + run_distributed(main, args) +``` + +- [x] **Step 4: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_entrypoint_uses_main.py -v` +Expected: PASS + +- [x] **Step 5: Manual end-to-end sanity check** + +`main()` has never actually executed for radar before — confirm it runs clean on a real (if synthetic) dataset before trusting the unit test alone. Build a small fixture (5 class dirs of CSVs under `/classes/`, balanced `annotations/instances_{train,test,val}_list.txt`) and drive `train.get_args_parser().parse_args([...]); train.run(args)` directly with `--device cpu`, a few epochs, `--quantization 0`. Confirm it completes and exports `model.onnx` without error. Then repeat with `--quantization 1` (previously silently a no-op under `main_debug`) and confirm the log now shows `QuantTrain` phase entries and a second exported model under the quantization output dir. + +- [x] **Step 6: Commit** + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/radar_classification/train.py tests/test_radar_entrypoint_uses_main.py +git commit -m "fix: radar_classification run() was dispatching to main_debug, not main" +``` + +--- + +## Task 2: Re-benchmark MPS vs CPU with the fix in place + +**Files:** +- None modified — this is a measurement task using the fixture/driver already built during investigation (`bench_radar.py`, `make_radar_fixture.py` in the session scratchpad — recreate if not available) + +**Interfaces:** +- Consumes: `radar_classification.train.run(args)` (now dispatching to `main`, per Task 1) + +- [x] **Step 1: Re-run the same benchmark used to characterize the bug** + +Using the same synthetic fixture (5 classes, balanced annotation lists, `LINEAR_4L_PC`, batch size 16, 30 epochs) and the same driver pattern that produced the pre-fix numbers (constructs `argv`, calls `train.get_args_parser().parse_args(argv)`, times `train.run(args)`), run once with `--device cpu` and once with `--device mps`. + +Run: `python bench_radar.py cpu` then `python bench_radar.py mps` (with `PYTORCH_ENABLE_MPS_FALLBACK=1` set) +Expected: both complete without error; timing printed as `total=...s for N epochs -> ...s/epoch` + +- [x] **Step 2: Record the comparison** + +Append a short results table to this plan (or a sibling `RESULTS.md` next to it) with pre-fix vs post-fix CPU and MPS per-epoch timing, and note whether `apply_hardware_defaults` actually enabled `torch.compile`/AMP for this run (check the training log for the relevant INFO lines from `compile_model_if_enabled`/`apply_hardware_defaults`). + +- [x] **Step 3: Commit** + +```bash +git add docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md +git commit -m "docs: record post-fix MPS vs CPU benchmark for radar training" +``` + +--- + +## Results + +Same fixture (5-class synthetic radar data), same driver (`bench_radar.py`, `LINEAR_4L_PC`, batch size 16, 30 epochs), run through `radar_classification.train.run(args)` — which now dispatches to `main()` per Task 1's fix. + +| Run | Device | Time/epoch | Total (30 epochs) | MPS vs CPU | +|---|---|---|---|---| +| Pre-fix (`main_debug`, historical) | CPU | 0.602s | — | — | +| Pre-fix (`main_debug`, historical) | MPS | 1.145s | — | **1.90x slower** | +| Post-fix (`main`) | CPU | 0.621s | 18.62s | — | +| Post-fix (`main`) | MPS | 1.066s | 31.99s | **1.72x slower** | + +**The gap did not close.** MPS is still substantially slower than CPU after the fix — 1.72x, essentially the same order as the pre-fix 1.90x. The ~9% wobble between the two ratios is consistent with ordinary run-to-run variance (different process, thermal state, etc.), not a structural change. + +**Why: `compile_model_if_enabled`/`apply_hardware_defaults` never engage for radar, in either `main()` or `main_debug()`.** Grepping both run logs (`radar_out_cpu/run.log`, `radar_out_mps/run.log`) for the exact strings `compile_model_if_enabled` emits on success (`"Compiling model with torch.compile"`), on quantization skip (`"compile_model is enabled but quantization is also enabled"`), and on failure (`"torch.compile failed"`) — zero matches in either log. Reading the code confirms why: + +- `radar_classification/train.py`'s `main()` (lines 187-334) calls `quantization_wrapped_model` (line 269) and `resume_from_checkpoint` (line 257) — those parts of Task 1's fix are real and now active — but it never imports or calls `compile_model_if_enabled` at all. Contrast with `timeseries_classification/train.py`, `timeseries_forecasting/train.py`, `timeseries_anomalydetection/train.py`, and `timeseries_regression/train.py`, which all import it from `..common.train_base` and call it before their training loops. `image_classification/train.py` and `audio_classification/train.py` also never call it — this looks like a scope gap in PR #22 ("compile-hardening", commit `dc8eeb2`), which only touched the four `timeseries_*` scripts, not the `radar_classification`/`image_classification`/`audio_classification` scripts. +- `apply_hardware_defaults` isn't part of `tinyml-tinyverse` at all — it lives in `tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py` and is wired only into `tinyml_modelmaker/ai_modules/timeseries/params.py` (commit `edbabba`, "wire apply_hardware_defaults into timeseries init_params"). It runs one layer up, in modelmaker's argv-construction step, before `train.py` ever sees `--compile-model`/`--native-amp`. `bench_radar.py` calls `radar_classification.train.get_args_parser().parse_args(argv)` directly, bypassing modelmaker entirely, so this layer was never in play for either the pre-fix or post-fix benchmark regardless. +- Net effect: `--compile-model` defaults to `0` and `--native-amp` defaults to `False` in `get_base_args_parser()` (`train_base.py:200,221`), nothing in the radar call path ever overrides them, and even if it did, `main()` has no code path that would act on them. torch.compile and AMP were off for all four runs (pre-fix CPU/MPS, post-fix CPU/MPS) — the comparison is apples-to-apples on that axis, just not for the reason the plan's premise expected. + +**Bottom line:** Task 1's fix is still correct and worth having — it wires up quantization (`QuantTrain` phase now actually runs quantization-aware training instead of a mislabeled second float pass) and checkpoint resume, both real, previously-silent gaps. But it does not touch torch.compile/AMP for radar, because that wiring was never built for the radar script in the first place (unlike the `timeseries_*` family). The MPS-slower-than-CPU result — expected for a small-op-heavy graph with no kernel fusion — persists post-fix because the mechanism that would fix it (`compile_model_if_enabled`) is not in `radar_classification/train.py`'s call graph at all. Closing this gap would require a separate change: wiring `compile_model_if_enabled` (and optionally AMP via `get_amp_context`/`get_grad_scaler`) into `radar_classification/train.py::main()`, mirroring what the `timeseries_*` scripts already do — out of scope for this plan. diff --git a/docs/superpowers/plans/2026-08-14-harden-compile-work-followups.md b/docs/superpowers/plans/2026-08-14-harden-compile-work-followups.md new file mode 100644 index 00000000..829b4379 --- /dev/null +++ b/docs/superpowers/plans/2026-08-14-harden-compile-work-followups.md @@ -0,0 +1,362 @@ +# Harden Compile-Hardening Follow-ups Implementation Plan + +**Goal:** Fix three issues an independent peer code review found in the compile-hardening plan's aftermath: a live crash reachable through a supported radar/timeseries config, three regression tests too weak to catch the failures they're meant to guard, and a DataLoader resource leak newly reachable now that `radar_classification/train.py`'s `main()` is the live entrypoint. + +**Architecture:** Three independent tasks, each touching different files with no shared state: +1. A one-condition change in two files (`radar_classification/train.py`, `timeseries_classification/train.py`) — fall back to loading datasets fresh when the quant-reuse cache is empty, instead of crashing. +2. Replace three source-text-grep tests with real behavioral tests that patch `compile_model_if_enabled` directly and assert on the call, mirroring an existing pattern already in this test suite (`test_anomalydetection_train_device_crash.py`). +3. Wrap radar's `main()` training body in `try/finally: shutdown_data_loaders(...)`, matching `image_classification/train.py`'s existing pattern exactly. + +**Tech Stack:** Python 3.10, PyTorch 2.7.1, pytest + +## Global Constraints + +- Python `==3.10.*` +- No new dependencies +- Task 1's fix must preserve the fast path (reuse cached dataset within the same process during a float-then-quant sequence) — only the "cache is empty" case should trigger a fresh load +- Task 2's new tests must genuinely fail if: the `compile_model_if_enabled` call is deleted, its return value is not assigned back to `model`, it's moved to the wrong position relative to `move_model_to_device`/`setup_distributed_model`, or it's placed in an unreachable branch +- Task 3 must not change any other structural aspect of radar's `main()` — pure re-indentation plus the try/finally wrapper and the one new call + +## Context: Why This Is Needed + +An independent peer code review of the compile-hardening plan (`docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md`) found three issues in code that plan touched or made newly reachable: + +**1. `run_quant_train_only: True` crashes for radar (and identically for `timeseries_classification`).** `main()`'s dataset loading (`radar_classification/train.py:193-200`) only calls `load_datasets(...)` when `args.quantization` is falsy; when true, it reads from the module-global `dataset_load_state` cache, populated only by a prior float-training call **in the same process**. `radar_base.py:401-426` has a documented, supported config (`run_quant_train_only`) that skips the float pass and calls `train_module.run(args)` once with `--quantization` set — so `dataset_load_state['dataset']` is still `None`, and `num_classes = len(dataset.classes)` (line 217) raises `AttributeError: 'NoneType' object has no attribute 'classes'`. Under `main_debug` (the function `main()` replaced) this config completed without crashing — `main_debug` loaded datasets unconditionally — but silently produced a non-quantized float model in the quantized output path, which is also wrong. The regression is "silently wrong" becoming "crashes loudly"; the fix here makes it neither. + +**2. Three of five new `compile_model_if_enabled` regression tests are source-text greps** (`tests/test_radar_compile_wired.py`, `tests/test_image_classification_compile_wired.py`, `tests/test_audio_classification_compile_wired.py`) — `inspect.getsource(main)` + `"compile_model_if_enabled(" in source`. This passes even if: the call's return value isn't assigned back to `model`; the call moves after `setup_distributed_model` or the `NeuralNetworkWithPreprocess` wrap (wrong shape); the call sits in a dead branch; or the string appears in a comment. This repo already has a stronger pattern for exactly this kind of assertion — `tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py`'s third test drives the real `main()` with its dependency chain mocked and asserts on what a specific downstream call received. `tinyml-tinyverse/tests/test_train_best_epoch_bugs_vision_audio.py` already has working, complete mock harnesses for `image_classification.train.main()` and `audio_classification.train.main()` specifically — reuse those as the base rather than building new ones. + +**3. `radar_classification/train.py`'s `main()` never calls `shutdown_data_loaders`.** `create_data_loaders` (`train_base.py:995-1000`) sets `persistent_workers=True` whenever `args.workers > 0` (default 8). `image_classification/train.py` (and all four `timeseries_*` scripts) wrap their training body in `try: ... finally: shutdown_data_loaders(data_loader, data_loader_test)` (see `image_classification/train.py:297,472-473`) specifically to clean these up. Radar's `main()` has no such wrapping at all. This was moot while `main_debug` (which built loaders without `persistent_workers`) was the live entrypoint; it's a real leak now that `main()` is live. + +--- + +## File Map + +| Action | Path | Responsibility | +|--------|------|-----------------| +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py:194` | Fall back to fresh load when quant-reuse cache is empty | +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py:214` | Same fix | +| Modify | `tinyml-tinyverse/tests/test_radar_compile_wired.py` | Replace source-grep with real behavioral test | +| Modify | `tinyml-tinyverse/tests/test_image_classification_compile_wired.py` | Same | +| Modify | `tinyml-tinyverse/tests/test_audio_classification_compile_wired.py` | Same | +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` | Wrap `main()`'s training body in `try/finally: shutdown_data_loaders(...)` | + +--- + +## Task 1: Fix the `run_quant_train_only` crash in radar and timeseries_classification + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` (lines 193-200) +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py` (lines 213-220) + +**Interfaces:** +- Consumes: `load_datasets` (already imported in both files, unchanged signature) + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: radar_classification.train.main() must not crash when +args.quantization is set and dataset_load_state's cache is empty (the +run_quant_train_only config: a standalone quantized run with no preceding +float-training call in the same process to populate the cache).""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import torch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +class _FakeDataset: + classes = ['a', 'b'] + inverse_label_map = {0: 'a', 1: 'b'} + X = torch.zeros((4, 8)) + Y = torch.zeros((4,), dtype=torch.long) + + def __getitem__(self, i): + return self.X[i], self.X[i], self.Y[i] + + def __len__(self): + return 4 + + +def test_main_does_not_crash_when_quantization_cache_is_empty(): + radar_train.dataset_load_state['dataset'] = None + radar_train.dataset_load_state['dataset_test'] = None + radar_train.dataset_load_state['train_sampler'] = None + radar_train.dataset_load_state['test_sampler'] = None + + args = Namespace( + quantization=True, data_path='/fake', output_dir='/tmp/fake-radar-quant-only', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=True, + model='LINEAR_4L_PC', model_config=None, model_spec=None, dual_op=False, + output_int=True, quantization_method='QAT', weight_bitwidth=8, + activation_bitwidth=8, epochs=1, start_epoch=0, label_smoothing=0.0, + distributed=False, apex=False, print_freq=10, opset_version=17, + gen_golden_vectors=False, DEBUG=False, + ) + + fake_dataset = _FakeDataset() + fake_loaders = ([1], [1]) + + with ExitStack() as stack: + stack.enter_context(patch.object( + radar_train, "setup_training_environment", + return_value=(radar_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(radar_train, "prepare_transforms")) + stack.enter_context(patch.object( + radar_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(radar_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(radar_train.models, "get_model", return_value=torch.nn.Linear(8, 2))) + stack.enter_context(patch.object(radar_train, "log_model_summary")) + stack.enter_context(patch.object(radar_train, "load_pretrained_weights", side_effect=lambda m, a, l: m)) + stack.enter_context(patch.object(radar_train, "handle_export_only", return_value=False)) + stack.enter_context(patch.object(radar_train, "move_model_to_device")) + stack.enter_context(patch.object(radar_train, "compile_model_if_enabled", side_effect=lambda m, a, l, **kw: m)) + stack.enter_context(patch.object( + radar_train, "setup_distributed_model", side_effect=lambda m, a, d: (m, m, None))) + stack.enter_context(patch.object( + radar_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(radar_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "quantization_wrapped_model", side_effect=lambda m, *a, **kw: m)) + stack.enter_context(patch.object(radar_train.utils, "train_one_epoch_classification")) + stack.enter_context(patch.object( + radar_train.utils, "evaluate_classification", + return_value=(1.0, 1.0, 1.0, {}, [], []))) + stack.enter_context(patch.object(radar_train, "save_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "save_on_master")) + stack.enter_context(patch.object(radar_train.utils, "print_file_level_classification_summary")) + stack.enter_context(patch.object(radar_train.utils, "export_model")) + stack.enter_context(patch.object(radar_train, "log_training_time")) + + # Should not raise. Pre-fix: AttributeError on dataset_load_state['dataset'] is None. + radar_train.main(0, args) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_quant_only_no_crash.py -v` +Expected: FAIL with `AttributeError: 'NoneType' object has no attribute 'classes'` + +- [x] **Step 3: Fix radar_classification/train.py** + +Change: +```python + if args.quantization: + dataset, dataset_test, train_sampler, test_sampler = (dataset_load_state['dataset'], dataset_load_state['dataset_test'], + dataset_load_state['train_sampler'], dataset_load_state['test_sampler']) + else: +``` +to: +```python + if args.quantization and dataset_load_state['dataset'] is not None: + dataset, dataset_test, train_sampler, test_sampler = (dataset_load_state['dataset'], dataset_load_state['dataset_test'], + dataset_load_state['train_sampler'], dataset_load_state['test_sampler']) + else: +``` + +- [x] **Step 4: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_quant_only_no_crash.py -v` +Expected: PASS + +- [x] **Step 5: Apply the identical fix to timeseries_classification/train.py** + +Same one-word condition change at line 214 (`if args.quantization:` → `if args.quantization and dataset_load_state['dataset'] is not None:`). Write an analogous test (`tests/test_timeseries_classification_quant_only_no_crash.py`), adapting the mock args/dataset shape to `timeseries_classification`'s `main()` signature (check its existing tests, e.g. any in `tests/` already driving this `main()`, for the right fake-dataset shape before writing from scratch). + +- [x] **Step 6: Run both new tests plus the full suite** + +Run: `cd tinyml-tinyverse && python -m pytest tests/ -v` +Expected: both new tests pass; the one pre-existing unrelated failure (`test_anomalydetection_train_device_crash.py::test_main_passes_a_torch_device_not_the_raw_args_device_string`, broken since PR #22, unrelated to this change) is the only failure, if any + +- [x] **Step 7: Commit** + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/radar_classification/train.py tinyml_tinyverse/references/timeseries_classification/train.py tests/test_radar_quant_only_no_crash.py tests/test_timeseries_classification_quant_only_no_crash.py +git commit -m "fix: run_quant_train_only crashed when the dataset-reuse cache was empty" +``` + +--- + +## Task 2: Strengthen the three weak compile_model_if_enabled tests + +**Files:** +- Modify: `tinyml-tinyverse/tests/test_radar_compile_wired.py` +- Modify: `tinyml-tinyverse/tests/test_image_classification_compile_wired.py` +- Modify: `tinyml-tinyverse/tests/test_audio_classification_compile_wired.py` + +**Interfaces:** +- Consumes: `tinyml_tinyverse.references.common.train_base.compile_model_if_enabled` (patch target), each module's own `main()` + +- [x] **Step 1: Read the two reference patterns before writing anything** + +`tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py` (specifically its third test, the one driving `anomaly_train.main(0, args)` with the full dependency chain mocked) is the template for "drive the real `main()`, assert on what a specific patched call received." `tinyml-tinyverse/tests/test_train_best_epoch_bugs_vision_audio.py` already contains complete, working mock harnesses for `image_classification.train.main()` and `audio_classification.train.main()` — reuse those fixtures/mocking setups rather than rebuilding them; only the specific assertion (on `compile_model_if_enabled`) is new. + +- [x] **Step 2: Rewrite test_radar_compile_wired.py** + +Replace the `inspect.getsource` substring check with a test that patches `compile_model_if_enabled` on the `radar_train` module, drives `main(0, args)` with everything else mocked (adapt the mock harness from Task 1's new radar test above, which already has a complete working mock set for `main()` — reuse it directly), and asserts: +- `compile_model_if_enabled` was called exactly once +- with `model` as the first positional arg (the pre-`setup_distributed_model`, pre-`NeuralNetworkWithPreprocess`-wrap model) +- with `input_shape=(1,) + dataset.X.shape[1:]` (assert the actual kwarg value, not just its presence) +- that its return value was what got passed into `setup_distributed_model` (i.e., the compiled/wrapped model is what continues through the rest of `main()`, not silently discarded) + +- [x] **Step 3: Run it, verify RED then GREEN** + +Temporarily verify this test fails against the *pre-Task-1-of-the-original-compile-hardening-plan* code shape (e.g. by checking it would fail if the `model = compile_model_if_enabled(...)` line were reverted to not exist, or not reassign `model`) — the simplest way is to comment out the assignment locally, confirm the new test fails, then restore it and confirm it passes. Document this RED/GREEN evidence in the report even though the underlying fix already shipped weeks ago; the point is proving the *new test* has teeth, not re-doing the original fix. + +- [x] **Step 4: Repeat Steps 2-3 for test_image_classification_compile_wired.py** + +Adapt using `test_train_best_epoch_bugs_vision_audio.py`'s existing `image_classification.train.main()` mock harness as the base. + +- [x] **Step 5: Repeat Steps 2-3 for test_audio_classification_compile_wired.py** + +Adapt using `test_train_best_epoch_bugs_vision_audio.py`'s existing `audio_classification.train.main()` mock harness as the base. + +- [x] **Step 6: Run the full suite** + +Run: `cd tinyml-tinyverse && python -m pytest tests/ -v` +Expected: all pass except the one known pre-existing unrelated failure + +- [x] **Step 7: Commit** + +```bash +cd tinyml-tinyverse +git add tests/test_radar_compile_wired.py tests/test_image_classification_compile_wired.py tests/test_audio_classification_compile_wired.py +git commit -m "test: replace source-text compile_model_if_enabled checks with real behavioral assertions" +``` + +--- + +## Task 3: Fix radar's missing shutdown_data_loaders (DataLoader worker leak) + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py` (imports ~line 65-89; `main()` body from the line after `create_data_loaders` at line 223 through the end of the function, currently ~line 336) + +**Interfaces:** +- Consumes: `shutdown_data_loaders` from `..common.train_base` (already used by `image_classification/train.py` and the four `timeseries_*` scripts — not yet imported in radar's file) + +- [x] **Step 1: Write a test that verifies shutdown_data_loaders is called** + +```python +"""Regression test: radar_classification.train.main() must call +shutdown_data_loaders() before returning, so persistent DataLoader worker +processes (enabled whenever --workers > 0, the default) are cleaned up. +image_classification and all four timeseries_* reference scripts already +do this; radar's main() did not.""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import torch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +class _FakeDataset: + classes = ['a', 'b'] + inverse_label_map = {0: 'a', 1: 'b'} + X = torch.zeros((4, 8)) + Y = torch.zeros((4,), dtype=torch.long) + + def __getitem__(self, i): + return self.X[i], self.X[i], self.Y[i] + + def __len__(self): + return 4 + + +def test_main_calls_shutdown_data_loaders(): + args = Namespace( + quantization=False, data_path='/fake', output_dir='/tmp/fake-radar-shutdown', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=True, + model='LINEAR_4L_PC', model_config=None, model_spec=None, dual_op=False, + output_int=True, quantization_method='QAT', weight_bitwidth=8, + activation_bitwidth=8, epochs=1, start_epoch=0, label_smoothing=0.0, + distributed=False, apex=False, print_freq=10, opset_version=17, + gen_golden_vectors=False, DEBUG=False, + ) + fake_dataset = _FakeDataset() + fake_loaders = ([1], [1]) + + with ExitStack() as stack: + stack.enter_context(patch.object( + radar_train, "setup_training_environment", + return_value=(radar_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(radar_train, "prepare_transforms")) + stack.enter_context(patch.object( + radar_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(radar_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(radar_train.models, "get_model", return_value=torch.nn.Linear(8, 2))) + stack.enter_context(patch.object(radar_train, "log_model_summary")) + stack.enter_context(patch.object(radar_train, "load_pretrained_weights", side_effect=lambda m, a, l: m)) + stack.enter_context(patch.object(radar_train, "handle_export_only", return_value=False)) + stack.enter_context(patch.object(radar_train, "move_model_to_device")) + stack.enter_context(patch.object(radar_train, "compile_model_if_enabled", side_effect=lambda m, a, l, **kw: m)) + stack.enter_context(patch.object( + radar_train, "setup_distributed_model", side_effect=lambda m, a, d: (m, m, None))) + stack.enter_context(patch.object( + radar_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(radar_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "quantization_wrapped_model", side_effect=lambda m, *a, **kw: m)) + stack.enter_context(patch.object(radar_train.utils, "train_one_epoch_classification")) + stack.enter_context(patch.object( + radar_train.utils, "evaluate_classification", + return_value=(1.0, 1.0, 1.0, {}, [], []))) + stack.enter_context(patch.object(radar_train, "save_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "save_on_master")) + stack.enter_context(patch.object(radar_train.utils, "print_file_level_classification_summary")) + stack.enter_context(patch.object(radar_train.utils, "export_model")) + stack.enter_context(patch.object(radar_train, "log_training_time")) + mock_shutdown = stack.enter_context(patch.object(radar_train, "shutdown_data_loaders")) + + radar_train.main(0, args) + + mock_shutdown.assert_called_once_with(*fake_loaders) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_shutdown_data_loaders.py -v` +Expected: FAIL — `shutdown_data_loaders` is never imported or called, so patching it and asserting `assert_called_once_with` fails (or the import patch itself fails since the name doesn't exist yet — either failure mode confirms the gap) + +- [x] **Step 3: Add the import** + +Add `shutdown_data_loaders` to the existing `from ..common.train_base import (...)` block in `radar_classification/train.py`. + +- [x] **Step 4: Wrap the training body in try/finally** + +Change the structure from (data loaders created, then everything else at the same indentation level through end of function) to: +```python + data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + + try: + logger.info("Creating model") + # ... (everything currently between here and the end of main(), re-indented one level deeper) + finally: + shutdown_data_loaders(data_loader, data_loader_test) +``` +Match `image_classification/train.py:297-473`'s exact wrapping boundaries and style — the `try:` opens right after `create_data_loaders`, the `finally:` closes right before the function ends (after the `gen_golden_vectors` block, which is the last statement in `main()`). + +- [x] **Step 5: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_radar_shutdown_data_loaders.py -v` +Expected: PASS + +- [x] **Step 6: Run the full suite plus a real (non-mocked) manual check** + +Run: `cd tinyml-tinyverse && python -m pytest tests/ -v` — expect only the known pre-existing unrelated failure. + +Then run an actual training pass on the synthetic radar fixture (reuse `make_radar_fixture.py`/the argv pattern from `docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md`'s Task 1) for a couple of epochs, `--device cpu`, to confirm the re-indented function still runs correctly end to end (the re-indentation is mechanical but touches ~110 lines — a real run is the best check that nothing was mis-indented into the wrong scope). + +- [x] **Step 7: Commit** + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/radar_classification/train.py tests/test_radar_shutdown_data_loaders.py +git commit -m "fix: radar_classification main() leaked persistent DataLoader workers" +``` diff --git a/docs/superpowers/plans/2026-08-14-remove-dead-sample-rate-flag.md b/docs/superpowers/plans/2026-08-14-remove-dead-sample-rate-flag.md new file mode 100644 index 00000000..a3badf9e --- /dev/null +++ b/docs/superpowers/plans/2026-08-14-remove-dead-sample-rate-flag.md @@ -0,0 +1,136 @@ +# Remove Dead --sample-rate Flag from audio_classification Implementation Plan + +**Goal:** Remove the `--sample-rate` CLI flag from `audio_classification`'s `train.py` and `test_onnx.py` — confirmed dead code that silently does nothing while looking like the thing that controls audio sample rate, trapping anyone who reads or invokes these scripts directly. + +**Architecture:** Two `parser.add_argument('--sample-rate', ...)` calls deleted (one per file, each script has its own independent `get_args_parser()`). No other code changes. `--sampling-rate` — a separate, shared flag defined once in `train_base.py`'s base parser, already `required=True`, already the only thing `GoogleSpeechCommandsDataset` actually reads, and already what `tinyml-modelmaker`'s production orchestration (`audio_base.py`) passes — continues unchanged as the sole real control for audio sample rate. Two regression tests confirm the dead flag is gone from each parser. + +**Tech Stack:** Python 3.10, PyTorch 2.7.1, pytest + +## Global Constraints + +- Python `==3.10.*` +- No new dependencies +- `--sampling-rate` (the flag that actually works, defined in `train_base.py`) must not be touched, renamed, or moved +- Removing `--sample-rate` must not change the parsed value of any other CLI argument + +## Context: Why This Is Needed + +Discovered during the `compile-hardening-radar-image-audio` plan's Task 3 (audio_classification), while building a synthetic benchmark fixture. Passing `--sampling-rate 1.0` (a value meaningful for the *other* reference scripts' generic FFT preprocessing, meaningless for audio) crashed `GoogleSpeechCommandsDataset`'s MFCC extraction with `RuntimeError: stft(...) : expected 0 < n_fft < 1, but got n_fft=0`, exposing the mechanism: + +- `train_base.py:113` defines the base parser's generic `--sampling-rate` (`type=float, required=True`, dest `sampling_rate`) — shared across all seven reference scripts, meant for FFT-domain preprocessing params like radar's/timeseries' frame sizing. +- `audio_classification/train.py:132` (and, confirmed separately, `test_onnx.py:69`) *also* defines its own `--sample-rate` (`default=16000, type=int`, dest `sample_rate`) — a second, differently-named, audio-specific-looking flag. +- `GoogleSpeechCommandsDataset.__init__` (`audio_dataset.py:143-146`) does `for key, value in kwargs.items(): setattr(self, key, ...)` over every parsed CLI arg, then reads `self.sampling_rate = int(getattr(self, "sampling_rate", 16000))` — this reads the *base parser's* `sampling_rate`, never `sample_rate`. Confirmed via repo-wide grep: no reference to `self.sample_rate` or `args.sample_rate` exists anywhere in `tinyml-tinyverse`, `tinyml-modelmaker`, or `tinyml-modelzoo`. +- `tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py:329,377` (the production orchestration layer) only ever passes `--sampling-rate` (from `self.params.data_processing_feature_extraction.sampling_rate`, default 16000) — it never passes `--sample-rate` at all. + +**Practical impact:** the modelmaker-driven production path is not currently broken by this — it happens to work because `--sampling-rate`'s default (16000) and `--sample-rate`'s default (16000) coincide, and modelmaker only ever sets the one that's actually read. The real risk is for anyone invoking these scripts directly (bypassing modelmaker, as this session's own benchmark driver did) or maintaining this code: `--sample-rate` reads as the obviously-correct flag to set for a non-default sample rate, silently does nothing, and the actually-effective flag (`--sampling-rate`) reads as generic/unrelated-to-audio. This is a maintainability trap, not an active data-corruption bug in the shipped product today. + +**Why removal, not redirection:** the alternative (making `GoogleSpeechCommandsDataset` read `sample_rate` instead of, or in preference to, `sampling_rate`) would change behavior for the working, already-relied-upon path (`--sampling-rate`, which `audio_base.py` already sets correctly) for no benefit, and risks introducing yet another dual-flag ambiguity. Since `--sample-rate` is provably read by nothing, deleting it is a pure simplification with no behavior change on any currently-working path. + +--- + +## File Map + +| Action | Path | Responsibility | +|--------|------|-----------------| +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py:132` | Remove dead `--sample-rate` argument | +| Modify | `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py:69` | Remove dead `--sample-rate` argument | +| Create | `tinyml-tinyverse/tests/test_audio_sample_rate_dead_flag_removed.py` | Regression test pinning both flags' removal | + +--- + +## Task 1: Remove `--sample-rate` from both audio_classification scripts, with regression test + +**Files:** +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py` (line 132) +- Modify: `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py` (line 69) +- Create: `tinyml-tinyverse/tests/test_audio_sample_rate_dead_flag_removed.py` + +**Interfaces:** +- Consumes: `tinyml_tinyverse.references.audio_classification.train.get_args_parser`, `tinyml_tinyverse.references.audio_classification.test_onnx.get_args_parser` (both already exist, no signature change) + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: audio_classification's train.py and test_onnx.py must NOT +define --sample-rate. It was dead code -- GoogleSpeechCommandsDataset +(audio_dataset.py) only ever reads self.sampling_rate, set from the shared +base parser's --sampling-rate (train_base.py:113), which tinyml-modelmaker's +production orchestration (audio_base.py) already passes correctly. Keeping +--sample-rate around silently does nothing while looking like the flag that +controls audio sample rate -- a maintainability trap for anyone reading or +invoking these scripts directly.""" +from tinyml_tinyverse.references.audio_classification import train, test_onnx + + +def _dests(parser): + return {action.dest for action in parser._actions} + + +def test_train_parser_has_no_dead_sample_rate_flag(): + dests = _dests(train.get_args_parser()) + assert "sample_rate" not in dests, ( + "train.py still defines --sample-rate, but GoogleSpeechCommandsDataset " + "never reads self.sample_rate -- only self.sampling_rate (from the " + "shared --sampling-rate flag). This dead flag should be removed." + ) + assert "sampling_rate" in dests, ( + "the real, working flag (--sampling-rate, from the shared base parser) " + "must still be present." + ) + + +def test_test_onnx_parser_has_no_dead_sample_rate_flag(): + dests = _dests(test_onnx.get_args_parser()) + assert "sample_rate" not in dests, ( + "test_onnx.py still defines --sample-rate, but GoogleSpeechCommandsDataset " + "never reads self.sample_rate -- only self.sampling_rate (from the " + "shared --sampling-rate flag). This dead flag should be removed." + ) + assert "sampling_rate" in dests, ( + "the real, working flag (--sampling-rate, from the shared base parser) " + "must still be present." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_audio_sample_rate_dead_flag_removed.py -v` +Expected: FAIL — both `test_..._has_no_dead_sample_rate_flag` tests fail, `"sample_rate"` is present in `dests` + +- [x] **Step 3: Remove the dead flag from train.py** + +In `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py`, delete line 132: + +```python + parser.add_argument('--sample-rate', help='Audio sample rate in Hz', default=16000, type=int) +``` + +- [x] **Step 4: Remove the dead flag from test_onnx.py** + +In `tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py`, delete line 69: + +```python + parser.add_argument('--sample-rate', help='Audio sample rate in Hz', default=16000, type=int) +``` + +- [x] **Step 5: Run test to verify it passes** + +Run: `cd tinyml-tinyverse && python -m pytest tests/test_audio_sample_rate_dead_flag_removed.py -v` +Expected: PASS + +- [x] **Step 6: Run the full existing test suite for this module to confirm no collateral breakage** + +Run: `cd tinyml-tinyverse && python -m pytest tests/ -v -k "audio"` +Expected: all pass (or only the same pre-existing failures already known from the compile-hardening plan's Task 3 report, if any — confirm via `git stash` on this task's own diff that any failure pre-exists and is unrelated, the same verification pattern used in that plan's Task 1-3 reports) + +- [x] **Step 7: Manual sanity check — confirm --sampling-rate alone still drives the real sample rate correctly** + +Reuse the synthetic audio fixture from the compile-hardening plan's Task 3 (`make_audio_fixture.py` in the session scratchpad — regenerate if not present) and drive `train.get_args_parser().parse_args([...]); train.run(args)` with `--sampling-rate 16000` (no `--sample-rate`, since it no longer exists), a couple of epochs, `--device cpu`. Confirm it completes cleanly and produces the same MFCC feature shape as before (spot-check `n_fft`/`hop_length` implied by the log or a quick shape print) — this is the same invocation this session's own benchmark driver already used successfully, now with one less (dead) flag in its argv. + +- [x] **Step 8: Commit** + +```bash +cd tinyml-tinyverse +git add tinyml_tinyverse/references/audio_classification/train.py tinyml_tinyverse/references/audio_classification/test_onnx.py tests/test_audio_sample_rate_dead_flag_removed.py +git commit -m "fix: remove dead --sample-rate flag from audio_classification scripts" +``` diff --git a/docs/superpowers/plans/2026-08-14-wire-compile-model-into-modelmaker.md b/docs/superpowers/plans/2026-08-14-wire-compile-model-into-modelmaker.md new file mode 100644 index 00000000..b614aa72 --- /dev/null +++ b/docs/superpowers/plans/2026-08-14-wire-compile-model-into-modelmaker.md @@ -0,0 +1,307 @@ +# Wire compile_model into modelmaker for Radar, Vision, and Audio Implementation Plan + +**Goal:** Make `--compile-model` reachable from the supported `tinyml-modelmaker` production path for `radar`, `vision` (image), and `audio` — closing the gap an independent peer review found in the `compile-hardening-radar-image-audio` plan: `compile_model_if_enabled` was wired into all three `tinyml-tinyverse` scripts, but none of the three modules' `tinyml-modelmaker` orchestration layers ever pass `--compile-model`, so the flag is permanently `0` (its default) for anyone using the actual product, not just direct script invocation. + +**Architecture:** `timeseries` already has this wiring end to end and is the template: a `compile_model=0` field in its `params.py` training dict, an `apply_hardware_defaults(params, user_training_keys)` call at the end of `init_params()` (which auto-flips `compile_model` to `1` when CUDA is available and the user hasn't explicitly set it), and a `'--compile-model', f'{getattr(self.params.training, "compile_model", 0)}'` entry in its argv builder. `apply_hardware_defaults` itself (`tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py`) already guards every field access with `hasattr`, and its own docstring says it was built for exactly this rollout ("hasattr guards keep this safe for params that don't carry these fields yet (vision, audio — Phase 2)") — so `apply_hardware_defaults` itself needs no changes, only the three modules' own params.py/argv-builder files. + +**Tech Stack:** Python 3.10, PyTorch 2.7.1, pytest + +## Global Constraints + +- Python `==3.10.*` +- No new dependencies +- Do not modify `apply_hardware_defaults` itself (`tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py`) — it already supports this rollout via its `hasattr` guards +- Do not modify `timeseries`'s existing wiring (reference only) +- `compile_model` must default to `0` in each module's params (matching timeseries) — `apply_hardware_defaults` is what conditionally raises it to `1`, not the static default +- Each task's fix must be verified against a real training run through the `tinyml-modelmaker` layer (not just a unit test of the params/argv construction), since this is exactly the integration layer that direct-script-only testing has already been shown (twice, in prior plans) to miss real bugs in + +## Context: Why This Is Needed + +`docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md` wired `compile_model_if_enabled` into `radar_classification`, `image_classification`, and `audio_classification`'s `main()` functions in `tinyml-tinyverse`, and its own stated Goal was "so all seven reference training scripts get the same hardware-acceleration path." An independent peer code review found this Goal was only half-achieved: `grep -rn compile_model tinyml-modelmaker/tinyml_modelmaker/ai_modules/{radar,vision,audio}/` returns nothing, versus `timeseries/params.py:160,232` and `timeseries_base.py:765`, which have the full three-piece wiring. Since `tinyml-modelmaker` is the actual product surface (the `tinyml-tinyverse` reference scripts are the layer it drives, not something end users invoke directly), the compile wiring from the prior plan currently does nothing for anyone using radar/vision/audio through the product. + +This plan closes that gap using the exact pattern `timeseries` and `apply_hardware_defaults` already establish — confirmed by direct inspection of all four modules' current source before writing this plan (not assumed by analogy). + +--- + +## File Map + +| Action | Path | Responsibility | +|--------|------|-----------------| +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py` | Add `compile_model=0` field + `apply_hardware_defaults` call | +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py` | Add `--compile-model` to train argv | +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py` | Same as radar | +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py` | Same as radar | +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py` | Same as radar | +| Modify | `tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py` | Same as radar | + +--- + +## Task 1: Wire compile_model into radar's modelmaker layer + +**Files:** +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py` (training dict at line 80; end of `init_params` at lines 198-199) +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py` (`_build_common_train_argv`, line 317) + +**Interfaces:** +- Consumes: `tinyml_modelmaker.utils.hardware_defaults.apply_hardware_defaults` (existing, unmodified, already imported this way in `timeseries/params.py:38`) + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: radar's modelmaker params must carry a compile_model +field, and apply_hardware_defaults must be invoked so it can auto-enable +compile on CUDA -- matching the pattern timeseries already has. Without +this, --compile-model is unreachable from the modelmaker (product) path +even though tinyml-tinyverse's radar_classification.train.main() already +supports it.""" +from tinyml_modelmaker.ai_modules.radar.params import init_params + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "radar's params.training has no compile_model field -- " + "apply_hardware_defaults can't act on it (it's hasattr-guarded), " + "and the field never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries) -- " + "apply_hardware_defaults is what conditionally raises it, not the static default." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-modelmaker && python -m pytest tests/test_radar_compile_model_param.py -v` +Expected: FAIL — `AttributeError` or `hasattr` returns `False`, `compile_model` not present + +- [x] **Step 3: Add the field to radar/params.py's training dict** + +In `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py`, in the `training=dict(...)` block, add `compile_model=0,` immediately before the `training_device='cuda', # 'cpu', 'cuda'` line: + +```python + momentum=0, + compile_model=0, # 1 to enable torch.compile (inductor on CUDA, aot_eager on MPS) + training_device='cuda', # 'cpu', 'cuda' +``` + +- [x] **Step 4: Add the apply_hardware_defaults call** + +Add the import near the top of the file (matching `timeseries/params.py:38`'s exact form): +```python +from ...utils.hardware_defaults import apply_hardware_defaults +``` + +Change the end of `init_params` from: +```python + params = utils.ConfigDict(default_params, *args, **kwargs) + return params +``` +to: +```python + user_training_keys = set(args[0].get('training', {}).keys()) \ + if args and isinstance(args[0], dict) else set() + params = utils.ConfigDict(default_params, *args, **kwargs) + apply_hardware_defaults(params, user_training_keys) + return params +``` +(matching `timeseries/params.py:227-232` exactly, including the comment above `user_training_keys` in that file explaining why the `isinstance` check exists — copy it verbatim for consistency.) + +- [x] **Step 5: Run test to verify it passes** + +Run: `cd tinyml-modelmaker && python -m pytest tests/test_radar_compile_model_param.py -v` +Expected: PASS + +- [x] **Step 6: Wire --compile-model into radar_base.py's train argv** + +In `tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py`'s `_build_common_train_argv`, change: +```python + '--distributed', f'{distributed}', + '--device', f'{device}', + + '--generic-model', f'{self.params.common.generic_model}', +``` +to: +```python + '--distributed', f'{distributed}', + '--device', f'{device}', + '--compile-model', f'{getattr(self.params.training, "compile_model", 0)}', + + '--generic-model', f'{self.params.common.generic_model}', +``` + +- [x] **Step 7: Write a test confirming the argv actually carries it** + +```python +"""Regression test: radar_base.py's train argv must include --compile-model, +sourced from params.training.compile_model -- otherwise the field added in +Step 3 has nowhere to go and remains dead.""" +from tinyml_modelmaker.ai_modules.radar.params import init_params +from tinyml_modelmaker.ai_modules.radar.training.tinyml_tinyverse.radar_base import BaseRadarModelTraining + + +def test_build_common_train_argv_includes_compile_model(): + params = init_params() + params.training.compile_model = 1 + + class _Dummy(BaseRadarModelTraining): + train_module = None + test_module = None + + instance = object.__new__(_Dummy) + instance.params = params + + argv = instance._build_common_train_argv(device="cpu", distributed=0) + assert "--compile-model" in argv, "argv builder never emits --compile-model" + idx = argv.index("--compile-model") + assert argv[idx + 1] == "1", ( + f"--compile-model should carry params.training.compile_model's value (1), got {argv[idx + 1]!r}" + ) +``` + +Run: `cd tinyml-modelmaker && python -m pytest tests/test_radar_compile_model_param.py -v` (add this test to the same file) +Expected: PASS after Step 6's change, would FAIL before it + +- [x] **Step 8: Manual end-to-end verification through the modelmaker layer** + +This is the layer prior plans' direct-script-only testing has already missed real bugs in twice — verify for real, not just via unit tests. Using the synthetic radar fixture (`make_radar_fixture.py` from the session scratchpad, or regenerate per `docs/superpowers/plans/2026-08-13-radar-training-entrypoint-fix.md`'s Task 1), drive radar training through `BaseRadarModelTraining.run()` (not directly through `tinyml_tinyverse.references.radar_classification.train`) with `params.training.compile_model = 1` set explicitly, `--device cpu`, a couple of epochs. Confirm the run log shows `compile_model_if_enabled`'s own `Compiling model with torch.compile` INFO line — proving the flag's value actually reaches the trainer through the full modelmaker → tinyverse argv-passing chain, not just that the argv list contains the right string. + +- [x] **Step 9: Commit** + +```bash +cd tinyml-modelmaker +git add tinyml_modelmaker/ai_modules/radar/params.py tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py tests/test_radar_compile_model_param.py +git commit -m "feat: wire compile_model into radar's modelmaker params/argv/hardware-defaults" +``` + +--- + +## Task 2: Wire compile_model into vision's (image_classification's) modelmaker layer + +**Files:** +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py` (training dict at line 80; end of `init_params` at lines 226-227) +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py` (`_build_common_train_argv`, line 337) + +**Interfaces:** +- Same as Task 1, applied to the vision module. Independent of Task 1 — same pattern, different files. + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: vision's modelmaker params must carry a compile_model +field, and apply_hardware_defaults must be invoked so it can auto-enable +compile on CUDA -- matching the pattern timeseries already has.""" +from tinyml_modelmaker.ai_modules.vision.params import init_params + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "vision's params.training has no compile_model field -- " + "apply_hardware_defaults can't act on it (it's hasattr-guarded), " + "and the field never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries) -- " + "apply_hardware_defaults is what conditionally raises it, not the static default." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-modelmaker && python -m pytest tests/test_vision_compile_model_param.py -v` +Expected: FAIL — `compile_model` not present on `params.training` + +- [x] **Step 3: Add `compile_model=0,` to vision/params.py's training dict**, immediately before `training_device=constants.TRAINING_DEVICE_CUDA,`. + +- [x] **Step 4: Add the `apply_hardware_defaults` import and call**, same exact pattern as Task 1 Step 4, applied to `vision/params.py`'s `init_params`. + +- [x] **Step 5: Run test to verify it passes** + +- [x] **Step 6: Wire `--compile-model` into image_base.py's train argv.** In `_build_common_train_argv` (line 337 area), insert after `'--device', f'{device}',` — same pattern as Task 1 Step 6, adapted to this file's exact surrounding lines (confirm the exact `--generic-model`/next-line context before editing, since `image_base.py` also has a `--sampling-rate` line here per the file map already inspected — insert `--compile-model` right after `--device`, before those). + +- [x] **Step 7: Write the argv-carries-it test**, mirroring Task 1 Step 7, targeting `ModelTraining`/`BaseModelTraining` in `image_base.py` (check the exact base class name in this file — it likely differs from radar's `BaseRadarModelTraining`, confirm before writing). + +- [x] **Step 8: Manual end-to-end verification.** Reuse the image fixture pattern from `docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md` Task 2 (`make_image_fixture.py`), drive training through vision's modelmaker `run()` (not directly through `tinyml_tinyverse.references.image_classification.train`) with `compile_model = 1` set, confirm the compile INFO log line appears. + +- [x] **Step 9: Commit** + +```bash +cd tinyml-modelmaker +git add tinyml_modelmaker/ai_modules/vision/params.py tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py tests/test_vision_compile_model_param.py +git commit -m "feat: wire compile_model into vision's modelmaker params/argv/hardware-defaults" +``` + +--- + +## Task 3: Wire compile_model into audio's modelmaker layer + +**Files:** +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py` (training dict at line 79; end of `init_params` at lines 213-214) +- Modify: `tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py` (`_build_common_train_argv`, line 317) + +**Interfaces:** +- Same as Tasks 1/2, applied to the audio module. Independent of both. + +- [x] **Step 1: Write the failing test** + +```python +"""Regression test: audio's modelmaker params must carry a compile_model +field, and apply_hardware_defaults must be invoked so it can auto-enable +compile on CUDA -- matching the pattern timeseries already has.""" +from tinyml_modelmaker.ai_modules.audio.params import init_params + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "audio's params.training has no compile_model field -- " + "apply_hardware_defaults can't act on it (it's hasattr-guarded), " + "and the field never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries) -- " + "apply_hardware_defaults is what conditionally raises it, not the static default." + ) +``` + +- [x] **Step 2: Run test to verify it fails** + +Run: `cd tinyml-modelmaker && python -m pytest tests/test_audio_compile_model_param.py -v` +Expected: FAIL — `compile_model` not present on `params.training` + +- [x] **Step 3: Add `compile_model=0,` to audio/params.py's training dict**, immediately before `training_device='cuda', # 'cpu', 'cuda'`. + +- [x] **Step 4: Add the `apply_hardware_defaults` import and call**, same pattern as Tasks 1/2. + +- [x] **Step 5: Run test to verify it passes** + +- [x] **Step 6: Wire `--compile-model` into audio_base.py's train argv.** In `_build_common_train_argv` (line 317 area), insert after `'--device', f'{device}',`, matching this file's exact surrounding structure (confirmed identical to radar's at this point per the earlier file map inspection — `'--generic-model'` follows). + +- [x] **Step 7: Write the argv-carries-it test**, mirroring Task 1 Step 7. + +- [x] **Step 8: Manual end-to-end verification.** Reuse the audio fixture pattern from `docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md` Task 3 (`make_audio_fixture.py`), drive training through audio's modelmaker `run()` with `compile_model = 1` set, confirm the compile INFO log line appears. Watch for the `--sampling-rate` requirement (`required=True` on the base parser) — audio_base.py already supplies it correctly (per the separately-fixed `--sample-rate` dead-flag plan), so this should just work, but confirm rather than assume. + +- [x] **Step 9: Commit** + +```bash +cd tinyml-modelmaker +git add tinyml_modelmaker/ai_modules/audio/params.py tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py tests/test_audio_compile_model_param.py +git commit -m "feat: wire compile_model into audio's modelmaker params/argv/hardware-defaults" +``` + +--- + +## Note for the final whole-plan review + +Once all three tasks land, re-verify `apply_hardware_defaults`'s CUDA-only auto-enable behavior doesn't regress anything on GX10: with all three modules now carrying `compile_model` fields, a modelmaker-driven run on GX10 (where `torch.cuda.is_available()` is `True`) will auto-set `compile_model=1` unless the user explicitly configured it — meaning `inductor` will now auto-engage for radar/vision/audio through the product path on any CUDA machine. Given this plan's own sibling plan found `inductor` to be a large *regression* at short (30-epoch) run lengths for these three models, this is worth flagging explicitly in the whole-plan review as a real behavior change on CUDA hardware, not just a "flag now works" story — cross-reference `docs/superpowers/plans/2026-08-13-compile-hardening-radar-image-audio.md`'s Addendum section for the actual numbers. + +**Decision (2026-08-14, repo owner, following the whole-plan review) — SUPERSEDED, see below:** keep auto-enable, matching `timeseries`'s existing behavior — no change to `apply_hardware_defaults` or its scope. Rationale: the sibling plan's regressions were measured on very short (30-epoch) synthetic benchmarks where one-time `torch.compile`/`inductor` warmup dominates the total; real training runs are almost certainly long enough to amortize that fixed cost and come out ahead, which is the same reasoning `apply_hardware_defaults` already applies to `timeseries`. Radar/vision/audio now behave consistently with `timeseries` rather than being a special case. If this assumption turns out to be wrong for real workloads (not just this session's tiny synthetic fixtures), revisit by either excluding these three modules from `apply_hardware_defaults`'s auto-enable or exposing `compile_model` in the params `descriptions` block so users can see and override it from the GUI (currently YAML-only, noted as a gap by the whole-plan review). + +**Decision (2026-08-14, repo owner, later the same day) — REVERSED:** the auto-enable above has been reverted for radar/vision/audio. `apply_hardware_defaults` is no longer called from these three modules' `init_params()`; `compile_model` is back to being a purely explicit opt-in field (default `0`), identical in effect to before this plan ran, except the field and its `--compile-model` argv wiring remain in place and reachable. + +Trigger: an independent analysis of `apply_hardware_defaults` (`ANALYSIS-cuda-auto-defaults.md`, written by a session working on a downstream consumer of this codebase) found and this session independently reproduced a correctness bug (its finding F-1): `apply_hardware_defaults` infers "what the user explicitly set" via `user_training_keys = set(args[0].get('training', {}).keys()) if args and isinstance(args[0], dict) else set()`. `ConfigDict`'s constructor also accepts a YAML file path string (not just a dict) as `args[0]` — a normal, supported call shape. When a caller passes a path, that expression silently evaluates to an empty set regardless of what the YAML actually contains, so *every* key in that YAML — including an explicit `compile_model: 0` or `native_amp: false` — reads as "not explicitly set" and gets silently overridden on any CUDA machine. This is a real, silent-override bug, not a theoretical one: the whole-plan review's own reasoning above (short synthetic benchmarks understate long-run compile amortization) is beside the point if the mechanism can flip a setting the user explicitly turned off. + +Rather than special-case the YAML-path branch inside `apply_hardware_defaults` under time pressure, the safer immediate fix is to stop calling it from radar/vision/audio (which never had this wiring before this plan) and leave `timeseries` — which predates this plan and whose existing behavior is out of scope here — untouched. This is a strict revert to pre-plan behavior for these three modules' hardware-default handling; the rest of this plan's deliverable (the `compile_model` field and argv wiring) is retained and unaffected. A regression test per module (`test_compile_model_not_auto_enabled_on_cuda`) now guards against silently re-wiring `apply_hardware_defaults` back in without addressing F-1 first. + +If `apply_hardware_defaults` is later fixed to correctly introspect YAML-path configs (not just dict configs) for their explicitly-set keys, re-auto-enabling for radar/vision/audio can be reconsidered — the reasoning in the superseded decision above about long-run compile amortization still stands on its own merits, independent of the bug that forced this reversal. diff --git a/tinyml-modelmaker/tests/test_audio_compile_model_param.py b/tinyml-modelmaker/tests/test_audio_compile_model_param.py new file mode 100644 index 00000000..24ab3fc1 --- /dev/null +++ b/tinyml-modelmaker/tests/test_audio_compile_model_param.py @@ -0,0 +1,70 @@ +"""Regression test: audio's modelmaker params must carry a compile_model +field, reachable via explicit user config -- matching the pattern +timeseries already has. Without this, --compile-model is unreachable from +the modelmaker (product) path even though tinyml-tinyverse's +audio_classification.train.main() already supports it. + +radar/vision/audio deliberately do NOT call apply_hardware_defaults: it +auto-raises compile_model to 1 on CUDA even when a caller passed an +explicit config (e.g. a YAML path, which apply_hardware_defaults can't +introspect for "user explicitly set" keys), silently overriding an +explicit compile_model=0. See docs/superpowers/plans/ +2026-08-14-wire-compile-model-into-modelmaker.md's Decision section.""" +from unittest.mock import patch + +from tinyml_modelmaker.ai_modules.audio.params import init_params +from tinyml_modelmaker.ai_modules.audio.training.tinyml_tinyverse.audio_base import BaseAudioModelTraining + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "audio's params.training has no compile_model field -- " + "it never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries)." + ) + + +def test_compile_model_not_auto_enabled_on_cuda(): + """Guards the revert: radar/vision/audio must NOT auto-raise + compile_model on CUDA (unlike timeseries). apply_hardware_defaults + can silently override an explicit compile_model=0 for YAML-path + configs it can't introspect -- see the module docstring above.""" + with patch('torch.cuda.is_available', return_value=True): + params = init_params() + assert params.training.compile_model == 0, ( + "compile_model changed even though nothing should auto-enable it " + "on CUDA for this module anymore -- did apply_hardware_defaults " + "get wired back in?" + ) + + +def test_build_common_train_argv_includes_compile_model(): + """Regression test: audio_base.py's train argv must include --compile-model, + sourced from params.training.compile_model -- otherwise the field added + above has nowhere to go and remains dead.""" + params = init_params() + params.training.compile_model = 1 + # dataset_path/data_dir default to None; _build_common_train_argv joins + # them with os.path.join (unlike other fields, not wrapped in an + # f-string), so they must be populated here the way the real pipeline + # populates them after dataset preparation, before this test can reach + # the --compile-model wiring under test. + params.dataset.dataset_path = "/tmp/fake_dataset" + params.dataset.data_dir = "data" + + class _Dummy(BaseAudioModelTraining): + train_module = None + test_module = None + + instance = object.__new__(_Dummy) + instance.params = params + + argv = instance._build_common_train_argv(device="cpu", distributed=0) + assert "--compile-model" in argv, "argv builder never emits --compile-model" + idx = argv.index("--compile-model") + assert argv[idx + 1] == "1", ( + f"--compile-model should carry params.training.compile_model's value (1), got {argv[idx + 1]!r}" + ) diff --git a/tinyml-modelmaker/tests/test_radar_compile_model_param.py b/tinyml-modelmaker/tests/test_radar_compile_model_param.py new file mode 100644 index 00000000..49ff3298 --- /dev/null +++ b/tinyml-modelmaker/tests/test_radar_compile_model_param.py @@ -0,0 +1,70 @@ +"""Regression test: radar's modelmaker params must carry a compile_model +field, reachable via explicit user config -- matching the pattern +timeseries already has. Without this, --compile-model is unreachable from +the modelmaker (product) path even though tinyml-tinyverse's +radar_classification.train.main() already supports it. + +radar/vision/audio deliberately do NOT call apply_hardware_defaults: it +auto-raises compile_model to 1 on CUDA even when a caller passed an +explicit config (e.g. a YAML path, which apply_hardware_defaults can't +introspect for "user explicitly set" keys), silently overriding an +explicit compile_model=0. See docs/superpowers/plans/ +2026-08-14-wire-compile-model-into-modelmaker.md's Decision section.""" +from unittest.mock import patch + +from tinyml_modelmaker.ai_modules.radar.params import init_params +from tinyml_modelmaker.ai_modules.radar.training.tinyml_tinyverse.radar_base import BaseRadarModelTraining + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "radar's params.training has no compile_model field -- " + "it never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries)." + ) + + +def test_compile_model_not_auto_enabled_on_cuda(): + """Guards the revert: radar/vision/audio must NOT auto-raise + compile_model on CUDA (unlike timeseries). apply_hardware_defaults + can silently override an explicit compile_model=0 for YAML-path + configs it can't introspect -- see the module docstring above.""" + with patch('torch.cuda.is_available', return_value=True): + params = init_params() + assert params.training.compile_model == 0, ( + "compile_model changed even though nothing should auto-enable it " + "on CUDA for this module anymore -- did apply_hardware_defaults " + "get wired back in?" + ) + + +def test_build_common_train_argv_includes_compile_model(): + """Regression test: radar_base.py's train argv must include --compile-model, + sourced from params.training.compile_model -- otherwise the field added + above has nowhere to go and remains dead.""" + params = init_params() + params.training.compile_model = 1 + # dataset_path/data_dir default to None; _build_common_train_argv joins + # them with os.path.join (unlike other fields, not wrapped in an + # f-string), so they must be populated here the way the real pipeline + # populates them after dataset preparation, before this test can reach + # the --compile-model wiring under test. + params.dataset.dataset_path = "/tmp/fake_dataset" + params.dataset.data_dir = "data" + + class _Dummy(BaseRadarModelTraining): + train_module = None + test_module = None + + instance = object.__new__(_Dummy) + instance.params = params + + argv = instance._build_common_train_argv(device="cpu", distributed=0) + assert "--compile-model" in argv, "argv builder never emits --compile-model" + idx = argv.index("--compile-model") + assert argv[idx + 1] == "1", ( + f"--compile-model should carry params.training.compile_model's value (1), got {argv[idx + 1]!r}" + ) diff --git a/tinyml-modelmaker/tests/test_radar_lr_scheduler_default.py b/tinyml-modelmaker/tests/test_radar_lr_scheduler_default.py new file mode 100644 index 00000000..9d413040 --- /dev/null +++ b/tinyml-modelmaker/tests/test_radar_lr_scheduler_default.py @@ -0,0 +1,28 @@ +"""Regression test: radar's modelmaker params default (lr_scheduler) must be +a value init_lr_scheduler actually accepts. Found via manual E2E verification +while wiring compile_model into radar's modelmaker layer -- any radar +training invoked through modelmaker with default params (no lr_scheduler +override) crashed during optimizer/scheduler setup with: +RuntimeError: Invalid lr scheduler 'constantlr'. Only StepLR, CosineAnnealingLR +and ExponentialLR are supported. +'constantlr' isn't one of those three, and isn't the special 'none' value +(which is what actually produces a genuinely constant LR via +ConstantLR(factor=1.0, total_iters=0)) -- it's an unsupported typo/misnaming.""" +import torch + +from tinyml_tinyverse.common.utils.utils import init_lr_scheduler + + +def test_radar_default_lr_scheduler_is_accepted_by_init_lr_scheduler(): + from tinyml_modelmaker.ai_modules.radar.params import init_params + + params = init_params() + lr_scheduler_value = params.training.lr_scheduler + + dummy_optimizer = torch.optim.SGD([torch.nn.Parameter(torch.zeros(1))], lr=0.01) + + # Should not raise. Pre-fix: RuntimeError("Invalid lr scheduler 'constantlr'...") + scheduler = init_lr_scheduler( + lr_scheduler=lr_scheduler_value, optimizer=dummy_optimizer, epochs=10, lr_warmup_epochs=0, + ) + assert scheduler is not None diff --git a/tinyml-modelmaker/tests/test_vision_compile_model_param.py b/tinyml-modelmaker/tests/test_vision_compile_model_param.py new file mode 100644 index 00000000..f8220386 --- /dev/null +++ b/tinyml-modelmaker/tests/test_vision_compile_model_param.py @@ -0,0 +1,70 @@ +"""Regression test: vision's modelmaker params must carry a compile_model +field, reachable via explicit user config -- matching the pattern +timeseries already has. Without this, --compile-model is unreachable from +the modelmaker (product) path even though tinyml-tinyverse's +image_classification.train.main() already supports it. + +radar/vision/audio deliberately do NOT call apply_hardware_defaults: it +auto-raises compile_model to 1 on CUDA even when a caller passed an +explicit config (e.g. a YAML path, which apply_hardware_defaults can't +introspect for "user explicitly set" keys), silently overriding an +explicit compile_model=0. See docs/superpowers/plans/ +2026-08-14-wire-compile-model-into-modelmaker.md's Decision section.""" +from unittest.mock import patch + +from tinyml_modelmaker.ai_modules.vision.params import init_params +from tinyml_modelmaker.ai_modules.vision.training.tinyml_tinyverse.image_base import BaseImageModelTraining + + +def test_init_params_carries_compile_model_field(): + params = init_params() + assert hasattr(params.training, "compile_model"), ( + "vision's params.training has no compile_model field -- " + "it never reaches the --compile-model argv flag." + ) + assert params.training.compile_model == 0, ( + "compile_model must default to 0 (matching timeseries)." + ) + + +def test_compile_model_not_auto_enabled_on_cuda(): + """Guards the revert: radar/vision/audio must NOT auto-raise + compile_model on CUDA (unlike timeseries). apply_hardware_defaults + can silently override an explicit compile_model=0 for YAML-path + configs it can't introspect -- see the module docstring above.""" + with patch('torch.cuda.is_available', return_value=True): + params = init_params() + assert params.training.compile_model == 0, ( + "compile_model changed even though nothing should auto-enable it " + "on CUDA for this module anymore -- did apply_hardware_defaults " + "get wired back in?" + ) + + +def test_build_common_train_argv_includes_compile_model(): + """Regression test: image_base.py's train argv must include --compile-model, + sourced from params.training.compile_model -- otherwise the field added + above has nowhere to go and remains dead.""" + params = init_params() + params.training.compile_model = 1 + # dataset_path/data_dir default to None; _build_common_train_argv joins + # them with os.path.join (unlike other fields, not wrapped in an + # f-string), so they must be populated here the way the real pipeline + # populates them after dataset preparation, before this test can reach + # the --compile-model wiring under test. + params.dataset.dataset_path = "/tmp/fake_dataset" + params.dataset.data_dir = "data" + + class _Dummy(BaseImageModelTraining): + train_module = None + test_module = None + + instance = object.__new__(_Dummy) + instance.params = params + + argv = instance._build_common_train_argv(device="cpu", distributed=0) + assert "--compile-model" in argv, "argv builder never emits --compile-model" + idx = argv.index("--compile-model") + assert argv[idx + 1] == "1", ( + f"--compile-model should carry params.training.compile_model's value (1), got {argv[idx + 1]!r}" + ) diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py index 584892ee..677b0840 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/params.py @@ -103,6 +103,7 @@ def init_params(*args, **kwargs): optimizer='sgd', weight_decay=1e-4, lr_scheduler='cosineannealinglr', + compile_model=0, # 1 to enable torch.compile (inductor on CUDA, aot_eager on MPS) training_device='cuda', # 'cpu', 'cuda' num_gpus=1, # 0,1 distributed=True, diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py index d62270e1..a7fab273 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/audio/training/tinyml_tinyverse/audio_base.py @@ -315,6 +315,7 @@ def _build_common_train_argv(self, device, distributed): '--lr-warmup-epochs', '1', '--distributed', f'{distributed}', '--device', f'{device}', + '--compile-model', f'{getattr(self.params.training, "compile_model", 0)}', '--generic-model', f'{self.params.common.generic_model}', diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py index d6717110..a49087df 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/params.py @@ -103,8 +103,9 @@ def init_params(*args, **kwargs): lambda_reg=0, optimizer='sgd', weight_decay=0, - lr_scheduler='constantlr', + lr_scheduler='none', # constant LR -- 'constantlr' is not a valid value for init_lr_scheduler momentum=0, + compile_model=0, # 1 to enable torch.compile (inductor on CUDA, aot_eager on MPS) training_device='cuda', # 'cpu', 'cuda' num_gpus=1, # 0,1 distributed=True, diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py index 28a0ec17..89c02039 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/radar/training/tinyml_tinyverse/radar_base.py @@ -315,6 +315,7 @@ def _build_common_train_argv(self, device, distributed): '--lr-warmup-epochs', '1', '--distributed', f'{distributed}', '--device', f'{device}', + '--compile-model', f'{getattr(self.params.training, "compile_model", 0)}', '--generic-model', f'{self.params.common.generic_model}', diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py index 34aba2a3..12c2be96 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/params.py @@ -104,6 +104,7 @@ def init_params(*args, **kwargs): optimizer='sgd', weight_decay=1e-4, lr_scheduler='cosineannealinglr', + compile_model=0, # 1 to enable torch.compile (inductor on CUDA, aot_eager on MPS) training_device=constants.TRAINING_DEVICE_CUDA, num_gpus=1, # 0,1 distributed=True, diff --git a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py index da110355..c918ff6a 100644 --- a/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py +++ b/tinyml-modelmaker/tinyml_modelmaker/ai_modules/vision/training/tinyml_tinyverse/image_base.py @@ -335,6 +335,7 @@ def _build_common_train_argv(self, device, distributed): '--lr-warmup-epochs', '1', '--distributed', f'{distributed}', '--device', f'{device}', + '--compile-model', f'{getattr(self.params.training, "compile_model", 0)}', '--generic-model', f'{self.params.common.generic_model}', '--sampling-rate', f'{self.params.data_processing_feature_extraction.sampling_rate}', diff --git a/tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py b/tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py index 2958e80f..1accdc68 100644 --- a/tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py +++ b/tinyml-tinyverse/tests/test_anomalydetection_train_device_crash.py @@ -130,7 +130,7 @@ def test_main_passes_a_torch_device_not_the_raw_args_device_string(): stack.enter_context(patch.object(anomaly_train, "load_pretrained_weights", side_effect=lambda model, a, l: model)) stack.enter_context(patch.object(anomaly_train, "handle_export_only", return_value=False)) stack.enter_context(patch.object(anomaly_train, "move_model_to_device")) - stack.enter_context(patch.object(anomaly_train, "compile_model_if_enabled", side_effect=lambda model, a, l: model)) + stack.enter_context(patch.object(anomaly_train, "compile_model_if_enabled", side_effect=lambda model, a, l, **kw: model)) stack.enter_context(patch.object(anomaly_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model)) stack.enter_context(patch.object(anomaly_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) stack.enter_context(patch.object( diff --git a/tinyml-tinyverse/tests/test_audio_classification_compile_wired.py b/tinyml-tinyverse/tests/test_audio_classification_compile_wired.py new file mode 100644 index 00000000..651d1e8b --- /dev/null +++ b/tinyml-tinyverse/tests/test_audio_classification_compile_wired.py @@ -0,0 +1,144 @@ +"""Regression test: audio_classification.train.main() must call +compile_model_if_enabled, matching the timeseries_* pattern (PR #22). +Without this, --compile-model is silently a no-op for audio classification +training regardless of what the caller requests. + +A prior version of this test only did an `inspect.getsource` substring +check for "compile_model_if_enabled(" in main()'s source. A peer review +found that pattern would pass under several real regressions: the call +deleted (but the string surviving in a comment), the call's return value +not reassigned to `model`, the call moved to the wrong position, or the +call placed in a dead/unreachable branch. This test instead patches +compile_model_if_enabled and drives the real main() (reusing the mock +harness from test_train_best_epoch_bugs_vision_audio.py), then asserts on +the actual call: it happens exactly once, with the pre-wrap model and the +expected input_shape kwarg, and its return value is what flows into +setup_distributed_model -- not silently discarded. +""" +import os +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from tinyml_tinyverse.references.audio_classification import train as audio_train + + +def _fake_data_loaders(): + item = (torch.tensor(0), torch.zeros(1, 3, 4), torch.tensor(0)) + return [item], [item] + + +def _base_audio_args(tmp_dir, **overrides): + tmp_dir = str(tmp_dir) + args = Namespace( + quantization=True, dont_train_just_feat_ext='False', load_saved_model='None', + nas_enabled=False, generic_model=False, nn_for_feature_extraction=False, + output_int=True, auto_quantization=False, distributed=False, gen_golden_vectors=False, + model='dummy', model_config=None, model_spec=None, dual_op=False, + label_smoothing=0.0, apex=False, print_freq=10, quantization_method='QAT', + epochs=3, start_epoch=3, # loop runs zero iterations, like a fully-resumed checkpoint + output_dir=tmp_dir, weight_bitwidth=8, activation_bitwidth=8, + autoquant_tolerance_classification=0.1, opset_version=17, device='cpu', + file_level_classification_log=os.path.join(tmp_dir, 'file_level.log'), DEBUG=False, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +class _FakeAudioClassificationDataset: + classes = ["a", "b"] + X = np.zeros((4, 3, 4), dtype=np.float32) + inverse_label_map = {0: "a", 1: "b"} + + +def _dataset_load_state_patch(train_module, dataset): + """patch.dict so the module-global dataset_load_state mutation is undone + after each test instead of leaking into later tests in the session.""" + return patch.dict(train_module.dataset_load_state, { + 'dataset': dataset, 'dataset_test': dataset, + 'train_sampler': None, 'test_sampler': None, + }) + + +def _patch_audio_common_pipeline(stack, dataset, pre_wrap_model, compiled_model_sentinel): + stack.enter_context(_dataset_load_state_patch(audio_train, dataset)) + stack.enter_context(patch.object( + audio_train, "setup_training_environment", + return_value=(audio_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(audio_train, "prepare_transforms")) + stack.enter_context(patch.object(audio_train, "create_data_loaders", return_value=_fake_data_loaders())) + stack.enter_context(patch.object(audio_train.models, "get_model", return_value=pre_wrap_model)) + stack.enter_context(patch.object(audio_train, "load_pretrained_weights", side_effect=lambda model, a, _logger: model)) + stack.enter_context(patch.object(audio_train, "handle_export_only", return_value=False)) + # Shared recorder so we can assert move_model_to_device runs BEFORE + # compile_model_if_enabled. move_model_to_device mutates the model + # in-place (model.to(device)) and has no captured return value in + # production code, so a bare no-op mock has no observable side effect + # that distinguishes call order -- side_effect here is what makes the + # order observable. + call_order = [] + stack.enter_context(patch.object( + audio_train, "move_model_to_device", + side_effect=lambda *a, **kw: call_order.append("move_model_to_device"))) + mock_compile = stack.enter_context(patch.object( + audio_train, "compile_model_if_enabled", + side_effect=lambda *a, **kw: (call_order.append("compile_model_if_enabled"), compiled_model_sentinel)[1])) + mock_setup_distributed = stack.enter_context(patch.object( + audio_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None))) + stack.enter_context(patch.object(audio_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(audio_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(audio_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model)) + stack.enter_context(patch.object(audio_train.utils, "export_model")) + stack.enter_context(patch.object(audio_train, "log_training_time")) + return mock_compile, mock_setup_distributed, call_order + + +def test_main_calls_compile_model_if_enabled_and_uses_its_return_value(tmp_path): + args = _base_audio_args(tmp_path) + dataset = _FakeAudioClassificationDataset() + + pre_wrap_model = torch.nn.Identity() + # Distinct object returned by compile_model_if_enabled, so we can prove it -- + # not the pre-wrap model -- is what continues into setup_distributed_model. + compiled_model_sentinel = torch.nn.Sequential(pre_wrap_model) + + with ExitStack() as stack: + mock_compile, mock_setup_distributed, call_order = _patch_audio_common_pipeline( + stack, dataset, pre_wrap_model, compiled_model_sentinel) + audio_train.main(0, args) + + # Called exactly once, with the pre-wrap model as the first positional arg. + mock_compile.assert_called_once() + call_args, call_kwargs = mock_compile.call_args + assert call_args[0] is pre_wrap_model, ( + "compile_model_if_enabled's first positional arg should be the model " + "returned by models.get_model() (pre-setup_distributed_model), but got " + "something else." + ) + + # input_shape kwarg must be the actual expected value, not merely present. + expected_input_shape = (1,) + dataset.X.shape[1:] + assert call_kwargs.get("input_shape") == expected_input_shape, ( + f"expected input_shape={expected_input_shape!r}, got {call_kwargs.get('input_shape')!r}" + ) + + # The compiled/wrapped return value must be what flows into + # setup_distributed_model -- not silently discarded. + mock_setup_distributed.assert_called_once() + setup_distributed_call_args = mock_setup_distributed.call_args[0] + assert setup_distributed_call_args[0] is compiled_model_sentinel, ( + "setup_distributed_model should receive compile_model_if_enabled's return " + "value, but received something else -- the compiled model's return value " + "looks like it was discarded (model = compile_model_if_enabled(...) not " + "assigned back to `model`)." + ) + + # Order matters: compile must run on a model that's already on its + # target device, not before. + assert call_order == ["move_model_to_device", "compile_model_if_enabled"], ( + f"Expected move_model_to_device to run before compile_model_if_enabled, got order: {call_order}" + ) diff --git a/tinyml-tinyverse/tests/test_audio_sample_rate_dead_flag_removed.py b/tinyml-tinyverse/tests/test_audio_sample_rate_dead_flag_removed.py new file mode 100644 index 00000000..c862da07 --- /dev/null +++ b/tinyml-tinyverse/tests/test_audio_sample_rate_dead_flag_removed.py @@ -0,0 +1,39 @@ +"""Regression test: audio_classification's train.py and test_onnx.py must NOT +define --sample-rate. It was dead code -- GoogleSpeechCommandsDataset +(audio_dataset.py) only ever reads self.sampling_rate, set from the shared +base parser's --sampling-rate (train_base.py:113), which tinyml-modelmaker's +production orchestration (audio_base.py) already passes correctly. Keeping +--sample-rate around silently does nothing while looking like the flag that +controls audio sample rate -- a maintainability trap for anyone reading or +invoking these scripts directly.""" +from tinyml_tinyverse.references.audio_classification import train, test_onnx + + +def _dests(parser): + return {action.dest for action in parser._actions} + + +def test_train_parser_has_no_dead_sample_rate_flag(): + dests = _dests(train.get_args_parser()) + assert "sample_rate" not in dests, ( + "train.py still defines --sample-rate, but GoogleSpeechCommandsDataset " + "never reads self.sample_rate -- only self.sampling_rate (from the " + "shared --sampling-rate flag). This dead flag should be removed." + ) + assert "sampling_rate" in dests, ( + "the real, working flag (--sampling-rate, from the shared base parser) " + "must still be present." + ) + + +def test_test_onnx_parser_has_no_dead_sample_rate_flag(): + dests = _dests(test_onnx.get_args_parser()) + assert "sample_rate" not in dests, ( + "test_onnx.py still defines --sample-rate, but GoogleSpeechCommandsDataset " + "never reads self.sample_rate -- only self.sampling_rate (from the " + "shared --sampling-rate flag). This dead flag should be removed." + ) + assert "sampling_rate" in dests, ( + "the real, working flag (--sampling-rate, from the shared base parser) " + "must still be present." + ) diff --git a/tinyml-tinyverse/tests/test_image_classification_compile_wired.py b/tinyml-tinyverse/tests/test_image_classification_compile_wired.py new file mode 100644 index 00000000..de72f4ad --- /dev/null +++ b/tinyml-tinyverse/tests/test_image_classification_compile_wired.py @@ -0,0 +1,139 @@ +"""Regression test: image_classification.train.main() must call +compile_model_if_enabled, matching the timeseries_* pattern (PR #22). +Without this, --compile-model is silently a no-op for image classification +training regardless of what the caller requests. + +A prior version of this test only did an `inspect.getsource` substring +check for "compile_model_if_enabled(" in main()'s source. A peer review +found that pattern would pass under several real regressions: the call +deleted (but the string surviving in a comment), the call's return value +not reassigned to `model`, the call moved to the wrong position, or the +call placed in a dead/unreachable branch. This test instead patches +compile_model_if_enabled and drives the real main() (reusing the mock +harness from test_train_best_epoch_bugs_vision_audio.py), then asserts on +the actual call: it happens exactly once, with the pre-wrap model and the +expected input_shape kwarg, and its return value is what flows into +setup_distributed_model -- not silently discarded. +""" +import os +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from tinyml_tinyverse.references.image_classification import train as img_train + + +def _fake_data_loaders(): + item = (torch.tensor(0), torch.zeros(1, 3, 4), torch.tensor(0)) + return [item], [item] + + +def _base_classification_args(tmp_dir, **overrides): + tmp_dir = str(tmp_dir) + args = Namespace( + quantization=True, dont_train_just_feat_ext='False', load_saved_model='None', + nas_enabled='False', generic_model=False, nn_for_feature_extraction=False, + output_int=True, auto_quantization=False, distributed=False, gen_golden_vectors=False, + model='dummy', model_config=None, model_spec=None, dual_op=False, + label_smoothing=0.0, apex=False, print_freq=10, quantization_method='QAT', + epochs=3, start_epoch=3, # loop runs zero iterations, like a fully-resumed checkpoint + output_dir=tmp_dir, weight_bitwidth=8, activation_bitwidth=8, + autoquant_tolerance_classification=0.1, opset_version=17, device='cpu', + file_level_classification_log=os.path.join(tmp_dir, 'file_level.log'), DEBUG=False, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +class _FakeImageClassificationDataset: + classes = ["a", "b"] + X = np.zeros((4, 3, 4), dtype=np.float32) + inverse_label_map = {0: "a", 1: "b"} + + +def _dataset_load_state_patch(train_module, dataset): + """patch.dict so the module-global dataset_load_state mutation is undone + after each test instead of leaking into later tests in the session.""" + return patch.dict(train_module.dataset_load_state, { + 'dataset': dataset, 'dataset_test': dataset, + 'train_sampler': None, 'test_sampler': None, + }) + + +def test_main_calls_compile_model_if_enabled_and_uses_its_return_value(tmp_path): + args = _base_classification_args(tmp_path) + dataset = _FakeImageClassificationDataset() + + pre_wrap_model = torch.nn.Identity() + # Distinct object returned by compile_model_if_enabled, so we can prove it -- + # not the pre-wrap model -- is what continues into setup_distributed_model. + compiled_model_sentinel = torch.nn.Sequential(pre_wrap_model) + + with ExitStack() as stack: + stack.enter_context(_dataset_load_state_patch(img_train, dataset)) + stack.enter_context(patch.object( + img_train, "setup_training_environment", + return_value=(img_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(img_train, "prepare_transforms")) + stack.enter_context(patch.object(img_train, "create_data_loaders", return_value=_fake_data_loaders())) + stack.enter_context(patch.object(img_train.models, "get_model", return_value=pre_wrap_model)) + stack.enter_context(patch.object(img_train, "load_pretrained_weights", side_effect=lambda model, a, _logger: model)) + stack.enter_context(patch.object(img_train, "handle_export_only", return_value=False)) + # Shared recorder so we can assert move_model_to_device runs BEFORE + # compile_model_if_enabled. move_model_to_device mutates the model + # in-place (model.to(device)) and has no captured return value in + # production code, so a bare no-op mock has no observable side + # effect that distinguishes call order -- side_effect here is what + # makes the order observable. + call_order = [] + stack.enter_context(patch.object( + img_train, "move_model_to_device", + side_effect=lambda *a, **kw: call_order.append("move_model_to_device"))) + mock_compile = stack.enter_context(patch.object( + img_train, "compile_model_if_enabled", + side_effect=lambda *a, **kw: (call_order.append("compile_model_if_enabled"), compiled_model_sentinel)[1])) + mock_setup_distributed = stack.enter_context(patch.object( + img_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None))) + stack.enter_context(patch.object(img_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(img_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(img_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model)) + stack.enter_context(patch.object(img_train.utils, "export_model")) + stack.enter_context(patch.object(img_train, "log_training_time")) + stack.enter_context(patch.object(img_train, "shutdown_data_loaders")) + img_train.main(0, args) + + # Called exactly once, with the pre-wrap model as the first positional arg. + mock_compile.assert_called_once() + call_args, call_kwargs = mock_compile.call_args + assert call_args[0] is pre_wrap_model, ( + "compile_model_if_enabled's first positional arg should be the model " + "returned by models.get_model() (pre-setup_distributed_model), but got " + "something else." + ) + + # input_shape kwarg must be the actual expected value, not merely present. + expected_input_shape = (1,) + dataset.X.shape[1:] + assert call_kwargs.get("input_shape") == expected_input_shape, ( + f"expected input_shape={expected_input_shape!r}, got {call_kwargs.get('input_shape')!r}" + ) + + # The compiled/wrapped return value must be what flows into + # setup_distributed_model -- not silently discarded. + mock_setup_distributed.assert_called_once() + setup_distributed_call_args = mock_setup_distributed.call_args[0] + assert setup_distributed_call_args[0] is compiled_model_sentinel, ( + "setup_distributed_model should receive compile_model_if_enabled's return " + "value, but received something else -- the compiled model's return value " + "looks like it was discarded (model = compile_model_if_enabled(...) not " + "assigned back to `model`)." + ) + + # Order matters: compile must run on a model that's already on its + # target device, not before. + assert call_order == ["move_model_to_device", "compile_model_if_enabled"], ( + f"Expected move_model_to_device to run before compile_model_if_enabled, got order: {call_order}" + ) diff --git a/tinyml-tinyverse/tests/test_radar_compile_wired.py b/tinyml-tinyverse/tests/test_radar_compile_wired.py new file mode 100644 index 00000000..5d0553ef --- /dev/null +++ b/tinyml-tinyverse/tests/test_radar_compile_wired.py @@ -0,0 +1,143 @@ +"""Regression test: radar_classification.train.main() must call +compile_model_if_enabled, matching the pattern already used in the 4 +timeseries_* reference scripts (PR #22). Without this, --compile-model is +silently a no-op for radar training regardless of what the caller requests. + +A prior version of this test only did an `inspect.getsource` substring +check for "compile_model_if_enabled(" in main()'s source. A peer review +found that pattern would pass under several real regressions: the call +deleted (but the string surviving in a comment), the call's return value +not reassigned to `model`, the call moved to the wrong position, or the +call placed in a dead/unreachable branch. This test instead patches +compile_model_if_enabled and drives the real main() (mocking everything +else, following the harness in test_radar_quant_only_no_crash.py), then +asserts on the actual call: it happens exactly once, with the pre-wrap +model and the expected input_shape kwarg, and its return value is what +flows into setup_distributed_model -- not silently discarded. +""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import torch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +class _FakeDataset: + classes = ['a', 'b'] + inverse_label_map = {0: 'a', 1: 'b'} + X = torch.zeros((4, 8)) + Y = torch.zeros((4,), dtype=torch.long) + + def __getitem__(self, i): + return self.X[i], self.X[i], self.Y[i] + + def __len__(self): + return 4 + + +def test_main_calls_compile_model_if_enabled_and_uses_its_return_value(): + radar_train.dataset_load_state['dataset'] = None + radar_train.dataset_load_state['dataset_test'] = None + radar_train.dataset_load_state['train_sampler'] = None + radar_train.dataset_load_state['test_sampler'] = None + + args = Namespace( + quantization=True, data_path='/fake', output_dir='/tmp/fake-radar-compile-wired', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=True, + model='LINEAR_4L_PC', model_config=None, model_spec=None, dual_op=False, + output_int=True, quantization_method='QAT', weight_bitwidth=8, + activation_bitwidth=8, epochs=1, start_epoch=0, label_smoothing=0.0, + distributed=False, apex=False, print_freq=10, opset_version=17, + gen_golden_vectors=False, DEBUG=False, + file_level_classification_log='/tmp/fake-radar-compile-wired/file_level.log', + ) + + fake_dataset = _FakeDataset() + # Each "batch" must be subscriptable like (raw, features, target) -- main()'s + # export step does `next(iter(data_loader_test))[1]` to get an example input. + fake_batch = (torch.zeros((1, 8)), torch.zeros((1, 8)), torch.zeros((1,), dtype=torch.long)) + fake_loaders = ([fake_batch], [fake_batch]) + + pre_wrap_model = torch.nn.Linear(8, 2) + # Distinct object returned by compile_model_if_enabled, so we can prove it -- + # not the pre-wrap model -- is what continues into setup_distributed_model. + compiled_model_sentinel = torch.nn.Sequential(pre_wrap_model) + + with ExitStack() as stack: + stack.enter_context(patch.object( + radar_train, "setup_training_environment", + return_value=(radar_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(radar_train, "prepare_transforms")) + stack.enter_context(patch.object( + radar_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(radar_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(radar_train.models, "get_model", return_value=pre_wrap_model)) + stack.enter_context(patch.object(radar_train, "log_model_summary")) + stack.enter_context(patch.object(radar_train, "load_pretrained_weights", side_effect=lambda m, a, l: m)) + stack.enter_context(patch.object(radar_train, "handle_export_only", return_value=False)) + # Shared recorder so we can assert move_model_to_device runs BEFORE + # compile_model_if_enabled. move_model_to_device mutates the model + # in-place (model.to(device)) and has no captured return value in + # production code, so a bare no-op mock has no observable side + # effect that distinguishes call order -- side_effect here is what + # makes the order observable. + call_order = [] + stack.enter_context(patch.object( + radar_train, "move_model_to_device", + side_effect=lambda *a, **kw: call_order.append("move_model_to_device"))) + mock_compile = stack.enter_context(patch.object( + radar_train, "compile_model_if_enabled", + side_effect=lambda *a, **kw: (call_order.append("compile_model_if_enabled"), compiled_model_sentinel)[1])) + mock_setup_distributed = stack.enter_context(patch.object( + radar_train, "setup_distributed_model", side_effect=lambda m, a, d: (m, m, None))) + stack.enter_context(patch.object( + radar_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(radar_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "quantization_wrapped_model", side_effect=lambda m, *a, **kw: m)) + stack.enter_context(patch.object(radar_train.utils, "train_one_epoch_classification")) + stack.enter_context(patch.object( + radar_train.utils, "evaluate_classification", + return_value=(1.0, 1.0, 1.0, {}, [], []))) + stack.enter_context(patch.object(radar_train, "save_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "save_on_master")) + stack.enter_context(patch.object(radar_train.utils, "print_file_level_classification_summary")) + stack.enter_context(patch.object(radar_train.utils, "export_model")) + stack.enter_context(patch.object(radar_train, "log_training_time")) + + radar_train.main(0, args) + + # Called exactly once, with the pre-wrap model as the first positional arg. + mock_compile.assert_called_once() + call_args, call_kwargs = mock_compile.call_args + assert call_args[0] is pre_wrap_model, ( + "compile_model_if_enabled's first positional arg should be the model " + "returned by models.get_model() (pre-setup_distributed_model, " + "pre-NeuralNetworkWithPreprocess-wrap), but got something else." + ) + + # input_shape kwarg must be the actual expected value, not merely present. + expected_input_shape = (1,) + fake_dataset.X.shape[1:] + assert call_kwargs.get("input_shape") == expected_input_shape, ( + f"expected input_shape={expected_input_shape!r}, got {call_kwargs.get('input_shape')!r}" + ) + + # The compiled/wrapped return value must be what flows into + # setup_distributed_model -- not silently discarded. + mock_setup_distributed.assert_called_once() + setup_distributed_call_args = mock_setup_distributed.call_args[0] + assert setup_distributed_call_args[0] is compiled_model_sentinel, ( + "setup_distributed_model should receive compile_model_if_enabled's return " + "value, but received something else -- the compiled model's return value " + "looks like it was discarded (model = compile_model_if_enabled(...) not " + "assigned back to `model`)." + ) + + # Order matters: compile must run on a model that's already on its + # target device, not before. + assert call_order == ["move_model_to_device", "compile_model_if_enabled"], ( + f"Expected move_model_to_device to run before compile_model_if_enabled, got order: {call_order}" + ) diff --git a/tinyml-tinyverse/tests/test_radar_entrypoint_uses_main.py b/tinyml-tinyverse/tests/test_radar_entrypoint_uses_main.py new file mode 100644 index 00000000..bf29425b --- /dev/null +++ b/tinyml-tinyverse/tests/test_radar_entrypoint_uses_main.py @@ -0,0 +1,25 @@ +"""Regression test: radar_classification.train.run() must dispatch to main(), +not main_debug() (a leftover notebook-parity harness that silently skips +quantization, AMP, and torch.compile for every radar training run).""" +from unittest.mock import patch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +def test_run_dispatches_to_main_not_main_debug(): + with patch.object(radar_train, "run_distributed") as mock_run_distributed: + fake_args = object() + radar_train.run(fake_args) + + mock_run_distributed.assert_called_once() + dispatched_fn, dispatched_args = mock_run_distributed.call_args[0] + + assert dispatched_fn is radar_train.main, ( + f"run() dispatched to {dispatched_fn.__name__!r}, expected 'main'. " + "main_debug() never applies quantization_wrapped_model or " + "resume_from_checkpoint -- wiring run() to it silently drops quantized " + "training and checkpoint resume for every radar run. main() is also the " + "only one of the two that can be extended with compile_model_if_enabled/" + "apply_hardware_defaults (added in a follow-up change)." + ) + assert dispatched_args is fake_args diff --git a/tinyml-tinyverse/tests/test_radar_quant_only_no_crash.py b/tinyml-tinyverse/tests/test_radar_quant_only_no_crash.py new file mode 100644 index 00000000..126c1709 --- /dev/null +++ b/tinyml-tinyverse/tests/test_radar_quant_only_no_crash.py @@ -0,0 +1,83 @@ +"""Regression test: radar_classification.train.main() must not crash when +args.quantization is set and dataset_load_state's cache is empty (the +run_quant_train_only config: a standalone quantized run with no preceding +float-training call in the same process to populate the cache).""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import torch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +class _FakeDataset: + classes = ['a', 'b'] + inverse_label_map = {0: 'a', 1: 'b'} + X = torch.zeros((4, 8)) + Y = torch.zeros((4,), dtype=torch.long) + + def __getitem__(self, i): + return self.X[i], self.X[i], self.Y[i] + + def __len__(self): + return 4 + + +def test_main_does_not_crash_when_quantization_cache_is_empty(): + radar_train.dataset_load_state['dataset'] = None + radar_train.dataset_load_state['dataset_test'] = None + radar_train.dataset_load_state['train_sampler'] = None + radar_train.dataset_load_state['test_sampler'] = None + + args = Namespace( + quantization=True, data_path='/fake', output_dir='/tmp/fake-radar-quant-only', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=True, + model='LINEAR_4L_PC', model_config=None, model_spec=None, dual_op=False, + output_int=True, quantization_method='QAT', weight_bitwidth=8, + activation_bitwidth=8, epochs=1, start_epoch=0, label_smoothing=0.0, + distributed=False, apex=False, print_freq=10, opset_version=17, + gen_golden_vectors=False, DEBUG=False, + file_level_classification_log='/tmp/fake-radar-quant-only/file_level.log', + ) + + fake_dataset = _FakeDataset() + # Each "batch" must be subscriptable like (raw, features, target) -- main()'s + # export step does `next(iter(data_loader_test))[1]` to get an example input. + fake_batch = (torch.zeros((1, 8)), torch.zeros((1, 8)), torch.zeros((1,), dtype=torch.long)) + fake_loaders = ([fake_batch], [fake_batch]) + + with ExitStack() as stack: + stack.enter_context(patch.object( + radar_train, "setup_training_environment", + return_value=(radar_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(radar_train, "prepare_transforms")) + stack.enter_context(patch.object( + radar_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(radar_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(radar_train.models, "get_model", return_value=torch.nn.Linear(8, 2))) + stack.enter_context(patch.object(radar_train, "log_model_summary")) + stack.enter_context(patch.object(radar_train, "load_pretrained_weights", side_effect=lambda m, a, l: m)) + stack.enter_context(patch.object(radar_train, "handle_export_only", return_value=False)) + stack.enter_context(patch.object(radar_train, "move_model_to_device")) + stack.enter_context(patch.object(radar_train, "compile_model_if_enabled", side_effect=lambda m, a, l, **kw: m)) + stack.enter_context(patch.object( + radar_train, "setup_distributed_model", side_effect=lambda m, a, d: (m, m, None))) + stack.enter_context(patch.object( + radar_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(radar_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "quantization_wrapped_model", side_effect=lambda m, *a, **kw: m)) + stack.enter_context(patch.object(radar_train.utils, "train_one_epoch_classification")) + stack.enter_context(patch.object( + radar_train.utils, "evaluate_classification", + return_value=(1.0, 1.0, 1.0, {}, [], []))) + stack.enter_context(patch.object(radar_train, "save_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "save_on_master")) + stack.enter_context(patch.object(radar_train.utils, "print_file_level_classification_summary")) + stack.enter_context(patch.object(radar_train.utils, "export_model")) + stack.enter_context(patch.object(radar_train, "log_training_time")) + + # Should not raise. Pre-fix: AttributeError on dataset_load_state['dataset'] is None. + radar_train.main(0, args) diff --git a/tinyml-tinyverse/tests/test_radar_shutdown_data_loaders.py b/tinyml-tinyverse/tests/test_radar_shutdown_data_loaders.py new file mode 100644 index 00000000..234c94ef --- /dev/null +++ b/tinyml-tinyverse/tests/test_radar_shutdown_data_loaders.py @@ -0,0 +1,85 @@ +"""Regression test: radar_classification.train.main() must call +shutdown_data_loaders() before returning, so persistent DataLoader worker +processes (enabled whenever --workers > 0, the default) are cleaned up. +image_classification and all four timeseries_* reference scripts already +do this; radar's main() did not.""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import torch + +from tinyml_tinyverse.references.radar_classification import train as radar_train + + +class _FakeDataset: + classes = ['a', 'b'] + inverse_label_map = {0: 'a', 1: 'b'} + X = torch.zeros((4, 8)) + Y = torch.zeros((4,), dtype=torch.long) + + def __getitem__(self, i): + return self.X[i], self.X[i], self.Y[i] + + def __len__(self): + return 4 + + +def test_main_calls_shutdown_data_loaders(): + radar_train.dataset_load_state['dataset'] = None + radar_train.dataset_load_state['dataset_test'] = None + radar_train.dataset_load_state['train_sampler'] = None + radar_train.dataset_load_state['test_sampler'] = None + + args = Namespace( + quantization=False, data_path='/fake', output_dir='/tmp/fake-radar-shutdown', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=True, + model='LINEAR_4L_PC', model_config=None, model_spec=None, dual_op=False, + output_int=True, quantization_method='QAT', weight_bitwidth=8, + activation_bitwidth=8, epochs=1, start_epoch=0, label_smoothing=0.0, + distributed=False, apex=False, print_freq=10, opset_version=17, + gen_golden_vectors=False, DEBUG=False, + file_level_classification_log='/tmp/fake-radar-shutdown/file_level.log', + ) + fake_dataset = _FakeDataset() + # Each "batch" must be subscriptable like (raw, features, target) -- main()'s + # export step does `next(iter(data_loader_test))[1]` to get an example input. + fake_batch = (torch.zeros((1, 8)), torch.zeros((1, 8)), torch.zeros((1,), dtype=torch.long)) + fake_loaders = ([fake_batch], [fake_batch]) + + with ExitStack() as stack: + stack.enter_context(patch.object( + radar_train, "setup_training_environment", + return_value=(radar_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(radar_train, "prepare_transforms")) + stack.enter_context(patch.object( + radar_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(radar_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(radar_train.models, "get_model", return_value=torch.nn.Linear(8, 2))) + stack.enter_context(patch.object(radar_train, "log_model_summary")) + stack.enter_context(patch.object(radar_train, "load_pretrained_weights", side_effect=lambda m, a, l: m)) + stack.enter_context(patch.object(radar_train, "handle_export_only", return_value=False)) + stack.enter_context(patch.object(radar_train, "move_model_to_device")) + stack.enter_context(patch.object(radar_train, "compile_model_if_enabled", side_effect=lambda m, a, l, **kw: m)) + stack.enter_context(patch.object( + radar_train, "setup_distributed_model", side_effect=lambda m, a, d: (m, m, None))) + stack.enter_context(patch.object( + radar_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(radar_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "quantization_wrapped_model", side_effect=lambda m, *a, **kw: m)) + stack.enter_context(patch.object(radar_train.utils, "train_one_epoch_classification")) + stack.enter_context(patch.object( + radar_train.utils, "evaluate_classification", + return_value=(1.0, 1.0, 1.0, {}, [], []))) + stack.enter_context(patch.object(radar_train, "save_checkpoint")) + stack.enter_context(patch.object(radar_train.utils, "save_on_master")) + stack.enter_context(patch.object(radar_train.utils, "print_file_level_classification_summary")) + stack.enter_context(patch.object(radar_train.utils, "export_model")) + stack.enter_context(patch.object(radar_train, "log_training_time")) + mock_shutdown = stack.enter_context(patch.object(radar_train, "shutdown_data_loaders")) + + radar_train.main(0, args) + + mock_shutdown.assert_called_once_with(*fake_loaders) diff --git a/tinyml-tinyverse/tests/test_timeseries_classification_quant_only_no_crash.py b/tinyml-tinyverse/tests/test_timeseries_classification_quant_only_no_crash.py new file mode 100644 index 00000000..2718e08b --- /dev/null +++ b/tinyml-tinyverse/tests/test_timeseries_classification_quant_only_no_crash.py @@ -0,0 +1,96 @@ +"""Regression test: timeseries_classification.train.main() must not crash when +args.quantization is set and dataset_load_state's cache is empty (the +run_quant_train_only config: a standalone quantized run with no preceding +float-training call in the same process to populate the cache). + +Identical bug and fix to radar_classification/train.py (see +tests/test_radar_quant_only_no_crash.py): pre-fix, `if args.quantization:` +always took the cache-reuse branch regardless of whether the cache had ever +been populated, so a lone quantized run crashed with +`AttributeError: 'NoneType' object has no attribute 'classes'` reading +`dataset.classes` from the never-populated `dataset_load_state['dataset']`. + +Fake dataset shape and mock pattern for driving this module's real main() +are adapted from tests/test_train_best_epoch_bugs_timeseries.py, which +already exercises tsc_train.main() end-to-end (zero-iteration-loop +"--resume" scenario) with quantization=True and a *populated* cache -- this +test instead drives it with an *empty* cache, so it must fall through to a +real (here, mocked) `load_datasets` call. +""" +from argparse import Namespace +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from tinyml_tinyverse.references.timeseries_classification import train as tsc_train + + +class _FakeTSClassificationDataset: + classes = ["a", "b"] + X = np.zeros((4, 3, 4), dtype=np.float32) + Y = np.zeros((4,), dtype=np.int64) + inverse_label_map = {0: "a", 1: "b"} + + +def _fake_data_loaders(): + item = (torch.tensor(0), torch.zeros(1, 3, 4), torch.tensor(0)) + return [item], [item] + + +def test_main_does_not_crash_when_quantization_cache_is_empty(): + tsc_train.dataset_load_state['dataset'] = None + tsc_train.dataset_load_state['dataset_test'] = None + tsc_train.dataset_load_state['train_sampler'] = None + tsc_train.dataset_load_state['test_sampler'] = None + + args = Namespace( + quantization=True, data_path='/fake', output_dir='/tmp/fake-tsc-quant-only', + gof_test=False, frame_size='None', dont_train_just_feat_ext='False', + load_saved_model='None', nas_enabled='False', generic_model=False, + nn_for_feature_extraction=False, model='dummy', model_config=None, + model_spec=None, dual_op=False, label_smoothing=0.0, apex=False, + print_freq=10, quantization_method='QAT', output_int=True, + auto_quantization=False, distributed=False, gen_golden_vectors=False, + epochs=3, start_epoch=3, # loop runs zero iterations -- avoids needing to + # mock train_one_epoch_classification/evaluate_classification/save_checkpoint, + # matching the existing zero-iteration-resume pattern in + # test_train_best_epoch_bugs_timeseries.py. + weight_bitwidth=8, activation_bitwidth=8, + autoquant_tolerance_classification=0.1, opset_version=17, device='cpu', + file_level_classification_log='/tmp/fake-tsc-quant-only/file_level.log', + DEBUG=False, + ) + + fake_dataset = _FakeTSClassificationDataset() + fake_loaders = _fake_data_loaders() + + with ExitStack() as stack: + stack.enter_context(patch.object( + tsc_train, "setup_training_environment", + return_value=(tsc_train.getLogger("test"), torch.device("cpu")))) + stack.enter_context(patch.object(tsc_train, "prepare_transforms")) + stack.enter_context(patch.object( + tsc_train, "load_datasets", + return_value=(fake_dataset, fake_dataset, None, None))) + stack.enter_context(patch.object(tsc_train.utils, "plot_feature_components_graph")) + stack.enter_context(patch.object(tsc_train, "create_data_loaders", return_value=fake_loaders)) + stack.enter_context(patch.object(tsc_train.models, "get_model", return_value=torch.nn.Identity())) + stack.enter_context(patch.object(tsc_train, "load_pretrained_weights", side_effect=lambda model, a, l: model)) + stack.enter_context(patch.object(tsc_train, "handle_export_only", return_value=False)) + stack.enter_context(patch.object(tsc_train, "move_model_to_device")) + stack.enter_context(patch.object(tsc_train, "compile_model_if_enabled", side_effect=lambda model, a, l, **kw: model)) + stack.enter_context(patch.object( + tsc_train, "setup_distributed_model", side_effect=lambda model, a, d: (model, model, None))) + stack.enter_context(patch.object(tsc_train, "setup_optimizer_and_scheduler", return_value=(MagicMock(), MagicMock()))) + stack.enter_context(patch.object(tsc_train, "resume_from_checkpoint")) + stack.enter_context(patch.object(tsc_train, "get_amp_context", return_value=MagicMock())) + stack.enter_context(patch.object(tsc_train, "get_grad_scaler", return_value=None)) + stack.enter_context(patch.object(tsc_train.utils, "quantization_wrapped_model", side_effect=lambda model, *a, **kw: model)) + stack.enter_context(patch.object(tsc_train.utils, "export_model")) + stack.enter_context(patch.object(tsc_train, "log_training_time")) + stack.enter_context(patch.object(tsc_train, "shutdown_data_loaders")) + + # Should not raise. Pre-fix: AttributeError on dataset_load_state['dataset'] is None. + tsc_train.main(0, args) diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py b/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py index cf2be6f3..679ddaf9 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/test_onnx.py @@ -64,9 +64,8 @@ def get_args_parser(): DESCRIPTION = "This script loads time series dataset and tests it against a onnx model using ONNX RT" parser = get_base_test_args_parser("This script loads an audio wav dataset and tests a classification model") - + # Audio preprocessing / feature extraction params - parser.add_argument('--sample-rate', help='Audio sample rate in Hz', default=16000, type=int) parser.add_argument('--audio-duration-ms', help='Audio clip duration in milliseconds', default=1000, type=int) parser.add_argument('--audio-feature', help='Audio feature type: MFCC, LPC, or RAW', default='MFCC', type=str) # MFCC params diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py index d7d3deb9..d6fe2909 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/audio_classification/train.py @@ -109,6 +109,7 @@ save_checkpoint, handle_export_only, move_model_to_device, + compile_model_if_enabled, log_training_time, apply_output_int_default, get_output_int_flag, @@ -128,7 +129,6 @@ def get_args_parser(): parser = get_base_args_parser("This script loads audio wav data and trains an classification model") - parser.add_argument('--sample-rate', help='Audio sample rate in Hz', default=16000, type=int) parser.add_argument('--audio-duration-ms', help='Audio clip duration in milliseconds', default=1000, type=int) parser.add_argument('--audio-feature', help='Audio feature type: MFCC, LPC, or RAW', default='MFCC', type=str) @@ -309,6 +309,7 @@ def main(gpu, args): return move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py index c154cd20..369a06ee 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/image_classification/train.py @@ -109,6 +109,7 @@ save_checkpoint, handle_export_only, move_model_to_device, + compile_model_if_enabled, log_training_time, apply_output_int_default, get_output_int_flag, @@ -327,6 +328,7 @@ def main(gpu, args): return move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py index 3fb0e9b9..e729c1d3 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/radar_classification/train.py @@ -82,10 +82,12 @@ save_checkpoint, handle_export_only, move_model_to_device, + compile_model_if_enabled, log_training_time, apply_output_int_default, get_output_int_flag, load_onnx_for_inference, + shutdown_data_loaders, ) dataset_loader_dict = {'GenericRadarDataset': GenericRadarDataset} @@ -190,7 +192,7 @@ def main(gpu, args): prepare_transforms(args) # Load or reuse datasets - if args.quantization: + if args.quantization and dataset_load_state['dataset'] is not None: dataset, dataset_test, train_sampler, test_sampler = (dataset_load_state['dataset'], dataset_load_state['dataset_test'], dataset_load_state['train_sampler'], dataset_load_state['test_sampler']) else: @@ -220,118 +222,122 @@ def main(gpu, args): logger.info("Loading data:") data_loader, data_loader_test = create_data_loaders(dataset, dataset_test, train_sampler, test_sampler, args, gpu) + try: - logger.info("Creating model") - if args.load_saved_model == 'None': - if args.nas_enabled == 'True': - if args.quantization: - model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + logger.info("Creating model") + if args.load_saved_model == 'None': + if args.nas_enabled == 'True': + if args.quantization: + model = torch.load(os.path.join(os.path.dirname(args.output_dir), os.path.join('base', 'nas_model.pt')), weights_only=False) + else: + nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) + model = search_and_get_model(nas_args) + if not model: + logger.error("Please check on prior errors. NAS wasn't able to create a model") + sys.exit(1) + torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) else: - nas_args = get_nas_args(args, data_loader, data_loader_test, num_classes, variables) - model = search_and_get_model(nas_args) - if not model: - logger.error("Please check on prior errors. NAS wasn't able to create a model") - sys.exit(1) - torch.save(model, os.path.join(args.output_dir, 'nas_model.pt')) + model = models.get_model( + args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, + model_spec=args.model_spec, + dual_op=args.dual_op) else: - model = models.get_model( - args.model, variables, num_classes, input_features=input_features, model_config=args.model_config, - model_spec=args.model_spec, - dual_op=args.dual_op) - else: - model = torch.load(args.load_saved_model, weights_only=False) + model = torch.load(args.load_saved_model, weights_only=False) - if args.generic_model or args.nas_enabled: - log_model_summary(model, args, variables, input_features, logger) + if args.generic_model or args.nas_enabled: + log_model_summary(model, args, variables, input_features, logger) - model = load_pretrained_weights(model, args, logger) + model = load_pretrained_weights(model, args, logger) - if handle_export_only(model, args, variables, input_features, logger): - return + if handle_export_only(model, args, variables, input_features, logger): + return - move_model_to_device(model, device, logger) - criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) + move_model_to_device(model, device, logger) + model = compile_model_if_enabled(model, args, logger, input_shape=(1,) + dataset.X.shape[1:]) + criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing) - model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) - optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) - resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) + model, model_without_ddp, model_ema = setup_distributed_model(model, args, device) + optimizer, lr_scheduler = setup_optimizer_and_scheduler(model, args) + resume_from_checkpoint(model_without_ddp, optimizer, lr_scheduler, model_ema, args) - phase = 'QuantTrain' if args.quantization else 'FloatTrain' - logger.info("Start training") - start_time = timeit.default_timer() - best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) + phase = 'QuantTrain' if args.quantization else 'FloatTrain' + logger.info("Start training") + start_time = timeit.default_timer() + best = dict(accuracy=0.0, f1=0, conf_matrix=dict(), epoch=None) - model = NeuralNetworkWithPreprocess(None, model) + model = NeuralNetworkWithPreprocess(None, model) - # if output_int not set by user, then set it to default of task_type - if args.output_int == None: - args.output_int = True - model = utils.quantization_wrapped_model( - model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, - args.epochs, args.output_int) - - + # if output_int not set by user, then set it to default of task_type + if args.output_int == None: + args.output_int = True + model = utils.quantization_wrapped_model( + model, args.quantization, args.quantization_method, args.weight_bitwidth, args.activation_bitwidth, + args.epochs, args.output_int) - for epoch in range(args.start_epoch, args.epochs): - if args.distributed: - train_sampler.set_epoch(epoch) - utils.train_one_epoch_classification( - model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, - print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, - is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) - if not (args.quantization_method in ['PTQ'] and args.quantization): - lr_scheduler.step() - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model, criterion, data_loader_test, device=device, transform=None, phase=phase, - num_classes=num_classes, dual_op=args.dual_op) - if model_ema: - avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( - model_ema, criterion, data_loader_test, device=device, transform=None, - log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) - if args.output_dir and avg_accuracy >= best['accuracy']: - logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") - best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch - best['predictions'], best['ground_truth'] = predictions, ground_truth - checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) - utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - # Log best epoch results - logger = getLogger(f"root.main.{phase}.BestEpoch") - logger.info("") - logger.info("Printing statistics of best epoch:") - logger.info(f"Best Epoch: {best['epoch']}") - logger.info(f"Acc@1 {best['accuracy']:.3f}") - logger.info(f"F1-Score {best['f1']:.3f}") - logger.info(f"AUC ROC Score {best['f1']:.3f}") - logger.info("") - logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], - columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], - index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), - headers="keys", tablefmt='grid'))) - Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, - name="root.utils.print_file_level_classification_summary", - append_log=True if args.quantization else False, console_log=False) - getLogger("root.utils.print_file_level_classification_summary").propagate = False - utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) - logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") - - # Export model - logger.info('Exporting model after training.') - if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): - example_input = next(iter(data_loader_test))[1] - input_shape = (1,) + dataset.X.shape[1:] - utils.export_model( - model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, - quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, - remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) - - log_training_time(start_time) + for epoch in range(args.start_epoch, args.epochs): + if args.distributed: + train_sampler.set_epoch(epoch) + utils.train_one_epoch_classification( + model, criterion, optimizer, data_loader, device, epoch, None, args.apex, model_ema, + print_freq=args.print_freq, phase=phase, num_classes=num_classes, dual_op=args.dual_op, + is_ptq=True if (args.quantization_method in ['PTQ'] and args.quantization) else False) + if not (args.quantization_method in ['PTQ'] and args.quantization): + lr_scheduler.step() + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model, criterion, data_loader_test, device=device, transform=None, phase=phase, + num_classes=num_classes, dual_op=args.dual_op) + if model_ema: + avg_accuracy, avg_f1, auc, avg_conf_matrix, predictions, ground_truth = utils.evaluate_classification( + model_ema, criterion, data_loader_test, device=device, transform=None, + log_suffix='EMA', print_freq=args.print_freq, phase=phase, dual_op=args.dual_op) + if args.output_dir and avg_accuracy >= best['accuracy']: + logger.info(f"Epoch {epoch}: {avg_accuracy:.2f} (Val accuracy) >= {best['accuracy']:.2f} (So far best accuracy). Hence updating checkpoint.pth") + best['accuracy'], best['f1'], best['auc'], best['conf_matrix'], best['epoch'] = avg_accuracy, avg_f1, auc, avg_conf_matrix, epoch + best['predictions'], best['ground_truth'] = predictions, ground_truth + checkpoint = save_checkpoint(model_without_ddp, optimizer, lr_scheduler, epoch, args, model_ema) + utils.save_on_master(checkpoint, os.path.join(args.output_dir, 'checkpoint.pth')) - if args.gen_golden_vectors: - generate_golden_vector_dir(args.output_dir) - output_int = get_output_int_flag(args) - generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model) + # Log best epoch results + logger = getLogger(f"root.main.{phase}.BestEpoch") + logger.info("") + logger.info("Printing statistics of best epoch:") + logger.info(f"Best Epoch: {best['epoch']}") + logger.info(f"Acc@1 {best['accuracy']:.3f}") + logger.info(f"F1-Score {best['f1']:.3f}") + logger.info(f"AUC ROC Score {best['f1']:.3f}") + logger.info("") + logger.info('Confusion Matrix:\n {}'.format(tabulate(pd.DataFrame(best['conf_matrix'], + columns=[f"Predicted as: {x}" for x in dataset.inverse_label_map.values()], + index=[f"Ground Truth: {x}" for x in dataset.inverse_label_map.values()]), + headers="keys", tablefmt='grid'))) + + Logger(log_file=args.file_level_classification_log, DEBUG=args.DEBUG, + name="root.utils.print_file_level_classification_summary", + append_log=True if args.quantization else False, console_log=False) + getLogger("root.utils.print_file_level_classification_summary").propagate = False + utils.print_file_level_classification_summary(dataset_test, best['predictions'], best['ground_truth'], phase) + logger.info(f"Generated file-level classification summary in: {args.file_level_classification_log}") + + # Export model + logger.info('Exporting model after training.') + if args.distributed is False or (args.distributed is True and int(os.environ['LOCAL_RANK']) == 0): + example_input = next(iter(data_loader_test))[1] + input_shape = (1,) + dataset.X.shape[1:] + utils.export_model( + model, input_shape=input_shape, output_dir=args.output_dir, opset_version=args.opset_version, + quantization=args.quantization, example_input=example_input, generic_model=args.generic_model, + remove_hooks_for_jit=True if (args.quantization_method == TinyMLQuantizationMethod.PTQ and args.quantization) else False) + + log_training_time(start_time) + + if args.gen_golden_vectors: + generate_golden_vector_dir(args.output_dir) + output_int = get_output_int_flag(args) + generate_golden_vectors(args.output_dir, dataset, output_int, args.generic_model) + finally: + shutdown_data_loaders(data_loader, data_loader_test) def main_debug(gpu, args): """Main training function for classification.""" @@ -521,7 +527,7 @@ def main_debug(gpu, args): def run(args): """Run training with optional distributed mode.""" - run_distributed(main_debug, args) + run_distributed(main, args) if __name__ == "__main__": diff --git a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py index b55e3707..007d3c29 100644 --- a/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py +++ b/tinyml-tinyverse/tinyml_tinyverse/references/timeseries_classification/train.py @@ -211,7 +211,7 @@ def main(gpu, args): prepare_transforms(args) # Load or reuse datasets - if args.quantization: + if args.quantization and dataset_load_state['dataset'] is not None: dataset, dataset_test, train_sampler, test_sampler = (dataset_load_state['dataset'], dataset_load_state['dataset_test'], dataset_load_state['train_sampler'], dataset_load_state['test_sampler']) else: