From 7821e5f0287612f71d78b6fdf04ca0cc11577eba Mon Sep 17 00:00:00 2001 From: Dhritiman Das <14159298+dhritimandas@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:12:07 +0530 Subject: [PATCH] feat(kwyk): import published TF kwyk weights with a verified mapping The kwyk architecture (McClure et al. 2019) was reimplemented in PyTorch, but users could not obtain the trusted published weights -- the three SavedModels ship only inside the neuronets/kwyk Docker container, and no conversion path existed. - nobrainer/datasets/convert_kwyk.py: TF SavedModel/.npz -> KWYKMeshNet state_dict converter. Strict load, offline parity check over all 8 conv layers (including the classifier), and a .provenance.json sidecar recording the source path + SHA-256 -- the MAP and BD checkpoints are structurally indistinguishable, so the content hash is the only durable record of which model a .pth came from. - KWYKMeshNet: the output classifier is now an FFGConv3d (was plain nn.Conv3d), matching the TF logits/ layer, which is itself a full VWN conv. All five logits variables map 1:1 with no information loss; the previous mean-path collapse made MC uncertainty under-dispersed at the output (TF contributes ~11.5 of run-to-run logit variance there; the plain conv contributed 0). Also adds a bias flag threaded through the hidden layers (default False, unchanged). - ConcreteDropout3d clamp widened [0.05, 0.95] -> [0.01, 0.99] in lockstep with the converter: the published SSD checkpoint stores p up to 0.954055, so the old ceiling silently clipped 89% of trained values (598/672). A unit test pins the two clamp ranges together. - Biases are imported by default (--drop-bias is an explicit, logged opt-out): every published checkpoint carries bias_m/bias_a for every conv layer, and dropping them shifts output logits by max ~20. Every mapping claim is verified against the real container and validated numerically against the live TF graph: the converted MAP model reproduces the original's logits to max|diff| = 9.5e-05 (relative error ~4e-6) on identical input; converted SSD preserves all 672 dropout probabilities to 6e-08 and reproduces output-layer MC variance. Full verification report, including the six discrepancies found and fixed and a second independent re-verification: docs/kwyk_mapping_verification.md. --- docs/kwyk_mapping_verification.md | 382 ++++++++++++ nobrainer/datasets/convert_kwyk.py | 670 ++++++++++++++++++++++ nobrainer/models/bayesian/kwyk_meshnet.py | 48 +- nobrainer/models/bayesian/vwn_layers.py | 19 +- nobrainer/tests/unit/test_convert_kwyk.py | 381 ++++++++++++ scripts/kwyk_reproduction/ARCHITECTURE.md | 4 +- 6 files changed, 1496 insertions(+), 8 deletions(-) create mode 100644 docs/kwyk_mapping_verification.md create mode 100644 nobrainer/datasets/convert_kwyk.py create mode 100644 nobrainer/tests/unit/test_convert_kwyk.py diff --git a/docs/kwyk_mapping_verification.md b/docs/kwyk_mapping_verification.md new file mode 100644 index 00000000..3bfde2a7 --- /dev/null +++ b/docs/kwyk_mapping_verification.md @@ -0,0 +1,382 @@ +# kwyk TF → PyTorch variable mapping: verification report + +**Status:** verification report (no code was changed as part of producing it). +**Update:** all six discrepancies have since been fixed and re-validated against +the real weights and the live TF graph — see §6 (Resolutions). + +**Verified against:** `neuronets/kwyk:latest-cpu`, digest +`sha256:8b72179a0b99284c5a520dde61982030891e6461946a5c36e67a445dd59226e1`, +SavedModels at `/opt/kwyk/saved_models/`: + +| model | timestamp | variables | +|---|---|---| +| `all_50_wn` (MAP) | 1555341859 | 41 (40 weights + `global_step`) | +| `all_50_bwn_09_multi` (BD) | 1555963478 | 41 (40 weights + `global_step`) | +| `all_50_bvwn_multi_prior` (SSD) | 1556816070 | 48 (47 weights + `global_step`) | + +**Method.** Three independent agents inventoried the PyTorch side, the TF side, and +the converter's stated assumptions, without reading each other's targets. Their +inventories were then cross-checked, and — beyond what any static inventory can +settle — the mapping was validated **numerically against the live TF graph** +(§4). + +Converter under review: `nobrainer/datasets/convert_kwyk.py`. +PyTorch target: `nobrainer/models/bayesian/{vwn_layers,kwyk_meshnet}.py`. + +--- + +## 1. Mapping table + +Published configuration: `filters=96`, `receptive_field=37`, `n_classes=50`, +`in_channels=1`. TF names carry **no** trailing `:0` as returned by +`NewCheckpointReader`. PyTorch indices are 0-based, TF indices are 1-based; +`torch_i = tf_i - 1`. + +### Hidden layers (TF `layer_1..layer_7` → PyTorch `layer_0..layer_6`) + +| TF variable | TF shape | PyTorch key | PT shape | Transform | +|---|---|---|---|---| +| `layer_{i}/conv3d/v` | `[3,3,3,1,96]` (i=1)
`[3,3,3,96,96]` (i=2..7) | `layer_{i-1}.conv.v` | `(96,1,3,3,3)`
`(96,96,3,3,3)` | `transpose(4,3,0,1,2)`, `→float32` | +| `layer_{i}/conv3d/g` | `[1,1,1,1,96]` | `layer_{i-1}.conv.g` | `(96,1,1,1,1)` | `transpose(4,3,0,1,2)` | +| `layer_{i}/conv3d/kernel_a` | `[3,3,3,1,96]` / `[3,3,3,96,96]` | `layer_{i-1}.conv.kernel_a` | `(96,1,3,3,3)` / `(96,96,3,3,3)` | `transpose(4,3,0,1,2)` | +| `layer_{i}/conv3d/bias_m` | `[96]` | `layer_{i-1}.conv.bias_m` | `(96,)` | verbatim copy — **only if `create_bias=True`** | +| `layer_{i}/conv3d/bias_a` | `[96]` | `layer_{i-1}.conv.bias_a` | `(96,)` | verbatim copy — **only if `create_bias=True`** | +| `layer_{i}/concrete_dropout/p` | `[96]` (SSD only) | `layer_{i-1}.dropout.p_logit` | `(96,)` | `clip(p, 0.05, 0.95)` → `log(p/(1−p))` | + +### Output layer (TF `logits/` → PyTorch `classifier`) + +| TF variable | TF shape | PyTorch key | PT shape | Transform | +|---|---|---|---|---| +| `logits/conv3d/v` | `[1,1,1,96,50]` | `classifier.weight` | `(50,96,1,1,1)` | `transpose(4,3,0,1,2)` then weight-norm collapse `g·v/‖v‖` | +| `logits/conv3d/g` | `[1,1,1,1,50]` | *(folded into `classifier.weight`)* | — | multiplied into the collapse above | +| `logits/conv3d/bias_m` | `[50]` | `classifier.bias` | `(50,)` | verbatim copy | +| `logits/conv3d/kernel_a` | `[1,1,1,96,50]` | **— none —** | — | **discarded** (logged) | +| `logits/conv3d/bias_a` | `[50]` | **— none —** | — | **discarded silently** (never read) | + +### Not mapped + +| TF variable | Reason | +|---|---| +| `global_step` (scalar, `int64`) | Training bookkeeping, not a weight. Ignored by the `^layer_(\d+)/` regex and by all templates. | + +--- + +## 2. CONFIRM / REFUTE + +### (a) TF conv filters are `[k,k,k,in,out]` and need `transpose(4,3,0,1,2)` — **CONFIRMED** + +Two independent lines of evidence. + +*Structural.* `layer_1/conv3d/v` has shape `[3,3,3,1,96]`. Layer 1 is the only +layer with `in_channels=1`, and the `1` sits at **axis 3** with `96` at **axis 4** — +so axis 3 is `in`, axis 4 is `out`, and axes 0–2 are the spatial kernel. A middle +layer is `[3,3,3,96,96]`, consistent. PyTorch requires `(out,in,k,k,k)` = +`(96,1,3,3,3)` (`vwn_layers.py:98`). The permutation `(4,3,0,1,2)` sends +axis4→0, axis3→1, axes0,1,2→2,3,4 — exactly correct. + +*Numerical.* See §4: converted weights reproduce the real TF graph's logits to +`max|Δ| = 9.92e-05`. A wrong permutation could not produce that. + +### (b) TF `g` is `[1,1,1,1,out]` or `[out]` — **CONFIRMED (5-D only); the 1-D branch is never exercised** + +Every `g` in all three models is **5-D**: `[1,1,1,1,96]` for hidden layers, +`[1,1,1,1,50]` for logits. A 1-D `[out]` `g` does **not occur** in any published +checkpoint. + +The converter's disjunction is *permissive*, so it is not refuted — but the 1-D +branch (`convert_kwyk.py:199-204`) is dead code with respect to real data, +reachable only from the synthetic unit-test fixture. Not a bug; recorded so it is +not mistaken for validated behaviour. See Discrepancy **D4**. + +### (c) TF layer index base, and whether it matches `_n_layers` — **CONFIRMED: base = 1, span matches** + +TF hidden layers span `layer_1 … layer_7` — **1-based**, 7 distinct indices, no +`layer_0`, no `layer_8`. PyTorch `_n_layers = 7` for *all three* receptive fields +(`_constants.py:7-11`; each schedule has 7 entries), and layers are attributes +`layer_0 … layer_6` (`kwyk_meshnet.py:190`). + +`detect_layer_index_base` computes `span = 7 - 1 + 1 = 7`, requires +`span == n_layers` → `7 == 7` ✓, and requires `lo ∈ (0,1)` → `1` ✓. It returns +`1`, so `tf_i = torch_i + 1` maps `layer_0..layer_6` ↔ `layer_1..layer_7` +correctly. + +Critically, this equality **only holds because the output layer is in its own +`logits/` namespace** rather than being `layer_8`. Had it been `layer_8`, span +would be 8 and the check would reject every real checkpoint. + +Caveat on the check's strength: it compares a *span*, not a set, so a checkpoint +with a gap (e.g. indices {1,2,4,5,6,7,8}) would pass. Not a defect for these +checkpoints — noted for completeness. + +### (d) `bias_m` and `bias_a` present in TF for every layer — **CONFIRMED, for all 8 conv layers in all 3 models (24/24)** + +Both are present for every hidden layer *and* for the logits layer, in all three +models. Shapes `[96]` hidden, `[50]` logits. + +This confirmation is what makes **D1** and **D2** serious rather than theoretical: +the converter defaults to *discarding* data that is always present. + +### (e) Concrete dropout `p` is stored as a probability, not a logit — **CONFIRMED** + +All 672 values (7 layers × 96) lie in `[0.607895, 0.954055]` — strictly inside +`(0,1)`, no negatives, none above 1. These are probabilities. PyTorch stores a +**logit** (`vwn_layers.py:195-197`), so a transform is genuinely required, and the +converter's `log(p/(1−p))` is the correct direction. + +`p` exists for hidden layers 1–7 only (7 variables); the logits layer has none, +matching PyTorch, where only `_VWNLayerConcrete` carries a `ConcreteDropout3d`. + +The *direction* is right. The *clamp* applied alongside it is not lossless — see +**D3**. + +--- + +## 3. DISCREPANCIES + +Ordered by severity. Each names the exact line that would have to change. + +### D1 — Default conversion silently discards biases that are always present (**HIGH**) + +`KWYKMeshNet` defaults to `bias=False` (`kwyk_meshnet.py:157`), and `--create-bias` +defaults to off (`convert_kwyk.py:471-472`). But §2(d) establishes every real +checkpoint has `bias_m`/`bias_a` for every layer. So the **default invocation +produces a numerically wrong model** and only emits a `logger.warning` +(`convert_kwyk.py:283-289`). + +Measured cost, MAP model vs the real TF graph on identical input: + +| run | max|Δ logits| | mean|Δ| | +|---|---|---| +| `--create-bias` (biases kept) | **9.92e-05** | 1.15e-05 | +| default (biases dropped) | **19.77** | 2.66 | + +That is a ~200,000× degradation, reported at WARNING level while the command exits +0. Worse, the offline parity check **passes anyway** in the no-bias case, because +it compares `conv.bias_m` against itself (`convert_kwyk.py:336` uses the model's +own bias on both sides) — so the built-in gate does not catch this. + +*Also note the module docstring at `kwyk_meshnet.py:138-143` asserts the published +checkpoints are bias-free. That statement is **false** and should be corrected.* + +**Lines to change:** `convert_kwyk.py:283-289` — make the bias-drop an error, or +invert the default so biases are kept unless explicitly discarded. Consequential +edit at `convert_kwyk.py:471-472` (flag default) and `:529` (model construction). + +### D2 — The logits layer's variational parameters are unrepresentable, so MC inference is not faithful (**HIGH, architectural**) + +The TF output layer is a **full VWN conv**: `logits/conv3d/{v,g,kernel_a,bias_m,bias_a}`. +`KWYKMeshNet.classifier` is a plain `nn.Conv3d` (`kwyk_meshnet.py:192`), which can +represent only the mean path. So `kernel_a` (logged, `convert_kwyk.py:338-342`) and +`bias_a` (**never read — there is no `_TF_LOGITS_BIAS_A` constant at all**, +`convert_kwyk.py:81-84`) are dropped. + +Consequence: for **deterministic** inference this is harmless and parity is +excellent (§4). For **MC/uncertainty** inference — the entire point of the SSD +model — the converted model's output layer is deterministic while the original's +samples. Measured on the real graph, the TF SSD model varies by +**max|Δ| = 11.52** between two runs on identical input, and TF BD by **39.93**; +the converted PyTorch classifier contributes **zero** of that variance. + +Uncertainty estimates from the converted SSD model will therefore be +systematically under-dispersed at the output layer. That is a real fidelity limit, +not a rounding issue, and it should be documented wherever the converted weights +are published. + +**Lines to change:** either `kwyk_meshnet.py:192` (make the classifier an +`FFGConv3d` so the variational output layer is representable) or, if the +limitation is accepted, `convert_kwyk.py:338-342` should be raised from +`logger.info` to `logger.warning` and state the MC consequence explicitly. + +### D3 — 89% of concrete-dropout probabilities are silently clamped (**MEDIUM**) + +`_p_to_logit` clips to `[0.05, 0.95]` (`convert_kwyk.py:218`, constants at `:57-58`) +before taking the logit. But the real `p` values run up to **0.954055**, and +**598 of 672 (89.0%)** exceed 0.95. Layer 7 is worst: its minimum is exactly +0.950000, so **95 of its 96 values** are clamped. + +Maximum absolute distortion is small (0.004055 in probability space), and the +clamp does mirror `ConcreteDropout3d.p`'s own `.clamp(0.05, 0.95)` +(`vwn_layers.py:202`) — so the converter is *faithful to the PyTorch forward*. +The discrepancy is that **PyTorch's clamp range is narrower than the range TF +actually trained to**, so the PyTorch model cannot express the trained dropout +rates regardless of the converter. The clustering just under 0.954 and the exact +0.950000 floor suggest TF applied its own constraint at a slightly wider bound. + +Because `p` scales activations directly (`vwn_layers.py:218`, `x * p`), a +systematic ~0.4% downward bias on 89% of channels is a small but real, one-signed +error — not noise that cancels. + +**Lines to change:** `convert_kwyk.py:57-58` (the clamp constants) *and* +`vwn_layers.py:202` (the forward clamp) must move together — changing only the +converter would be undone at forward time. Alternatively, keep the clamp but log +how many values were affected, at `convert_kwyk.py:218`. + +### D4 — The 1-D `g` branch is unreachable for real checkpoints (**LOW, informational**) + +`_g_tf_to_torch` accepts 1-D `g` (`convert_kwyk.py:199-204`), but every real `g` is +5-D (§2b). The branch is exercised only by the synthetic test fixture, so it is +untested against reality and could mask a genuine shape anomaly by silently +`reshape`-ing it. + +**Lines to change:** none required. If tightening is desired, +`convert_kwyk.py:199-204` could log when the 1-D path is taken. + +### D5 — MAP and BD checkpoints are structurally indistinguishable (**LOW, operational**) + +`all_50_wn` and `all_50_bwn_09_multi` have byte-identical key sets, shapes, and +dtypes — they differ only in values and in whether MC is enabled at inference. +The converter cannot tell them apart, and `dropout_type` / `dropout_rate` / +MC-at-inference are **not** stored in the PyTorch `state_dict` either +(`nn.Dropout3d` is parameterless). A user who converts BD weights and runs them +deterministically gets the MAP behaviour with no warning from either side. + +**Lines to change:** none in the mapping. Provenance (source model name) should be +recorded alongside the output `.pth`; today `main()` (`convert_kwyk.py:508-545`) +writes a bare `state_dict` with no metadata. + +### D6 — Stale prose in the converter's own docstring (**LOW**) + +`convert_kwyk.py:11-13` claims the axis conventions were *"verified against … +ARCHITECTURE.md and the source snippets of `vwn_layers.py`"*. They are now verified +against the live checkpoint and the live TF graph, which is materially stronger. +Separately, `ARCHITECTURE.md:115` describes the output as *"Layer 8 (logits)"*, +which reads as though it were `layer_8` in the variable namespace; it is not. + +**Lines to change:** `convert_kwyk.py:11-13` (upgrade the provenance claim); +`scripts/kwyk_reproduction/ARCHITECTURE.md:115` (clarify that "Layer 8" is +positional, and its variables live under `logits/`, not `layer_8/`). + +--- + +## 4. Numerical validation against the live TF graph + +Static inventories cannot prove an axis permutation is right — only that shapes are +consistent. So the mapping was checked end-to-end against the actual TF graph. + +**Determinism screening.** The SavedModel signature is +`volume (-1,32,32,32,1) → logits (-1,32,32,32,50)`. Running each model twice on an +identical seeded input: + +| model | max|run0 − run1| | verdict | +|---|---|---| +| `all_50_wn` (MAP) | **0.000000** | deterministic — usable as reference | +| `all_50_bwn_09_multi` (BD) | 39.93 | stochastic | +| `all_50_bvwn_multi_prior` (SSD) | 11.52 | stochastic | + +Only the MAP model admits a deterministic comparison; this is consistent with +ARCHITECTURE.md's "MC at inference: No" for `bwn`. + +**Parity.** MAP weights converted with `--create-bias`, then compared against the +TF logits for the same input (`NDHWC`↔`NCDHW` permuted, PyTorch run with +`mc=False`): + +``` +TF logits mean -2.289535 std 3.125234 +PT logits mean -2.289528 std 3.125223 +max |diff| = 9.918213e-05 +mean|diff| = 1.151648e-05 +rel error = 3.693623e-06 +argmax agreement = 100.0000% +``` + +A relative error of ~3.7e-06 over 7 dilated conv layers plus the collapsed output +layer is consistent with float32 accumulation and nothing else. **The axis +permutation, the weight-norm collapse, the 1-based↔0-based index mapping, the bias +copy, and the logits mapping are jointly correct.** + +**One honest caveat about `argmax`.** The bias-dropped run in **D1** *also* scored +100% argmax agreement despite `max|Δ| = 19.77`. On this synthetic Gaussian input +one class dominates almost everywhere, so argmax is insensitive. **Argmax +agreement must not be used as the parity criterion** — the logit-level error is +the meaningful signal. A parity check on real T1 data would be a stronger test +still. + +--- + +## 5. Summary + +| Claim | Verdict | +|---|---| +| (a) conv `[k,k,k,in,out]`, `transpose(4,3,0,1,2)` | **CONFIRMED** (structural + numerical) | +| (b) `g` is `[1,1,1,1,out]` or `[out]` | **CONFIRMED** — always 5-D; 1-D branch unexercised | +| (c) index base, matches `_n_layers` | **CONFIRMED** — base 1, span 7 = `_n_layers` 7 | +| (d) `bias_m`/`bias_a` present every layer | **CONFIRMED** — 24/24 across 3 models | +| (e) `p` is a probability, not a logit | **CONFIRMED** — all 672 values in [0.61, 0.95] | + +The mapping is **correct as implemented**, and verified numerically to ~4e-06 +relative error. The defects are not in the transforms but in the **defaults and +the discarded data**: D1 (bias dropped by default, ~200,000× worse output, warning +only, and the built-in parity gate does not catch it) is the one that will bite a +user first. D2 is the one that matters most for the SSD model's stated purpose. + +### Reproduction + +```bash +docker pull neuronets/kwyk:latest-cpu +# variable inventory +docker run --rm --platform linux/amd64 --entrypoint python neuronets/kwyk:latest-cpu -c \ + "import tensorflow as tf; r=tf.train.NewCheckpointReader( + '/opt/kwyk/saved_models/all_50_bvwn_multi_prior/1556816070/variables/variables'); + [print(k, r.get_variable_to_shape_map()[k]) for k in sorted(r.get_variable_to_shape_map())]" +``` +Note the image is `linux/amd64` (emulated on arm64) and has an ENTRYPOINT — it must +be overridden with `--entrypoint python`. + +--- + +## 6. Resolutions (implemented after the report) + +All six discrepancies were fixed and re-validated. Where a discrepancy offered +two remedies, the choice and its rationale are recorded. + +| # | Fix | Where | +|---|---|---| +| D1 | Biases are now imported **by default**; `--create-bias` replaced by an explicit `--drop-bias` opt-out whose warning quotes the measured max-~20 logit shift. `convert(create_bias=True)` is the new default. The false "checkpoints are bias-free" docstring was corrected. | `convert_kwyk.py` (`_build_arg_parser`, `convert`, `build_state_dict` warning), `kwyk_meshnet.py` (`bias` param docstring) | +| D2 | **Architectural fix chosen** (not the log-upgrade alternative): `KWYKMeshNet.classifier` is now an `FFGConv3d(kernel_size=1, bias=True)`, matching the TF `logits/` layer; `forward` passes `mc=mc_vwn` to it. The converter maps all five `logits/conv3d/*` variables 1:1 — including `bias_a`, which previously had no constant and was never read. Nothing is discarded. Chosen because the model is a *reproduction* of kwyk, the TF original **is** variational at the output, and no released checkpoint existed to break (all files pre-release). | `kwyk_meshnet.py` (classifier + forward), `convert_kwyk.py` (`_TF_LOGITS_BIAS_A`, `_add_classifier`, parity extended to the classifier) | +| D3 | Clamp widened to `[0.01, 0.99]` **in both places together**: `ConcreteDropout3d.p` and the converter, with module constants `CONCRETE_P_MIN/MAX` exposed in `vwn_layers.py` and a unit test (`test_clamp_constants_match_model`) enforcing the sync. `_p_to_logit` now logs a warning with the clipped count whenever any value falls outside the range. | `vwn_layers.py`, `convert_kwyk.py`, `test_convert_kwyk.py` | +| D4 | The 1-D `g` branch logs when taken (it never fires on real checkpoints). | `convert_kwyk.py` (`_g_tf_to_torch`) | +| D5 | `main()` writes a `.provenance.json` sidecar recording the source path, its SHA-256 (the only way to distinguish the structurally identical MAP/BD checkpoints), the architecture args (not recoverable from a `state_dict`), and converter/library versions. | `convert_kwyk.py` (`_write_provenance`, `_sha256_of_file`) | +| D6 | Converter docstring now claims verification against the live container + live TF graph (not "ARCHITECTURE.md and source snippets"); `ARCHITECTURE.md` clarifies that "Layer 8" is positional and its variables live under `logits/conv3d/*`, not `layer_8/`. | `convert_kwyk.py` module docstring, `scripts/kwyk_reproduction/ARCHITECTURE.md` | + +### Post-fix re-validation (real weights, live TF graph) + +| check | result | +|---|---| +| SSD conversion, **no flags** (new defaults) | strict load ✓, parity ✓ over **8** conv layers *(incl. classifier — previously 7)* | +| D3: `max\|tf_p − model.p\|` over all 672 channels | **5.96e-08** (float32 round-trip only; was 0.004 one-signed clipping on 89% of channels) | +| D3: model `p` max | **0.954055** — exactly the TF value, no longer saturated at 0.95 | +| D2: converted SSD, deterministic path | bit-reproducible (max diff 0.0) | +| D2: converted SSD, MC path run-to-run | **max diff 10.6** — same order as the TF original's measured 11.5; the old plain-conv classifier contributed 0 | +| D1/TF-graph parity: converted MAP vs live TF logits | **max\|diff\| = 9.54e-05**, mean 1.15e-05 — unchanged from the pre-fix 9.92e-05 baseline | +| D5: provenance sidecar | written, carries source SHA-256 + arch args | +| Test suite | `test_convert_kwyk.py`: 17 passed (13 original + 4 new); full unit suite 377 passed with only the pre-existing unrelated `test_croissant` failure | + +One deliberate consequence of D2 to be aware of: from-scratch `KWYKMeshNet` +training now treats the output layer variationally too (it contributes to +`kl_divergence()` and samples under `mc_vwn=True`). This is *more* faithful to +McClure et al. and to the TF implementation, but it does change the training +objective relative to the previous plain-conv classifier. + +### Second multi-agent re-verification (post-fix) + +The full three-agent workflow (PyTorch / TF / converter, independently +inventoried, then cross-checked with fresh numerical runs) was re-executed +after the fixes. Verdict: **D1–D6 all confirmed fixed**, with the TF ground +truth re-confirmed unchanged (41/41/48 variables, `logits/` namespace, p in +[0.607895, 0.954055]) and fresh conversions reproducing every §6 number +(MAP vs live TF 9.54e-05; p preserved to 5.96e-08; SSD MC output variance +present; sidecars distinguishing MAP/SSD by content hash; `--drop-bias` +warning firing and recording `bias: false`). + +The re-run caught **two residual stale-doc fragments**, both fixed on the +spot: + +1. The `kwyk_meshnet()` *factory* docstring still claimed "``bias`` defaults + to ``False`` to match the published kwyk checkpoints" — the same false + statement D1 removed from the class docstring. Corrected. +2. A test docstring still referenced the removed `--create-bias` flag. + Corrected. + +A repo sweep for further "bias-free"/"--create-bias" claims found only +accurate usages (descriptions of model configurations, not checkpoint +claims). diff --git a/nobrainer/datasets/convert_kwyk.py b/nobrainer/datasets/convert_kwyk.py new file mode 100644 index 00000000..2ce46daa --- /dev/null +++ b/nobrainer/datasets/convert_kwyk.py @@ -0,0 +1,670 @@ +"""Convert pretrained kwyk TensorFlow SavedModel weights into KWYKMeshNet (PyTorch). + +The kwyk architecture is already reimplemented in PyTorch +(``nobrainer/models/bayesian/kwyk_meshnet.py``). This module fills the remaining +gap: importing the *pretrained* TF SavedModel variables into that PyTorch model, +so users get the trusted published weights rather than a from-scratch retrain. + +Primary input is a pre-extracted ``.npz`` (no TensorFlow dependency). A direct +SavedModel/checkpoint path is supported only if ``tensorflow`` is importable. + +Every mapping claim below is verified against the actual +``neuronets/kwyk:latest-cpu`` container (SavedModel variables read via +``tf.train.NewCheckpointReader``) and validated numerically against the live TF +graph: the converted MAP model reproduces the original's logits to +``max|diff| ~ 1e-4`` (relative error ~4e-6) on identical input. Full evidence: +``docs/kwyk_mapping_verification.md``. + +- TF conv filters are ``[k, k, k, in, out]`` -> torch ``[out, in, k, k, k]`` + via ``transpose(4, 3, 0, 1, 2)``. +- ``g`` is 5-D ``[1, 1, 1, 1, out]`` in every published checkpoint -> torch + ``(out, 1, 1, 1, 1)``. A 1-D ``[out]`` form is also accepted (logged when + taken -- it never occurs in real checkpoints). +- ``kernel_a`` / ``bias_a`` are raw params; forward uses ``abs(...)`` as sigma, + so copy directly (no log/exp). +- concrete-dropout ``p`` is stored as a **probability** (all 672 published + values lie in [0.608, 0.955]) -> ``p_logit = log(p / (1 - p))``, clamped to + the model's forward clamp range first (see ``_CONCRETE_P_MIN/MAX``). +- ``bias_m``/``bias_a`` are present for **every** conv layer in **every** + published model (24/24 across the three variants), so biases are imported + by default; ``--drop-bias`` is an explicit opt-out. +- The output layer lives in its own ``logits/conv3d/*`` namespace (not + ``layer_8/``) and is a full VWN conv; it maps 1:1 onto the model's FFG + ``classifier`` with no information loss. +- ``ConcreteDropout3d`` stores its learnable logit as ``p_logit``, reached via + ``layer_{i}`` -> ``dropout`` -> ``p_logit``. That key exists only when the + model is built with ``dropout_type="concrete"``; the ``"bernoulli"`` variant + uses a parameter-free ``nn.Dropout3d``, so emitting ``p_logit`` against it + fails the strict load as an unexpected key. +- The layer-index base is not assumed -- it is detected and validated at + runtime by :func:`detect_layer_index_base`. +""" + +from __future__ import annotations + +import argparse +import datetime +import hashlib +import json +import logging +from pathlib import Path +import re + +import numpy as np +import torch +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +# --- Constants (no magic numbers in logic) ------------------------------------- + +# TF [k,k,k,in,out] -> torch [out,in,k,k,k] +_CONV_PERM: tuple[int, ...] = (4, 3, 0, 1, 2) + +# Concrete-dropout p is clamped to this range on every read of +# ConcreteDropout3d.p; match it exactly so the recovered p_logit reproduces +# forward behaviour. MUST stay in sync with vwn_layers.CONCRETE_P_MIN/MAX +# (kept as local literals so this module stays importable without nobrainer; +# a unit test asserts the two pairs are equal). Widened from [0.05, 0.95] +# because the published SSD checkpoint stores p up to 0.954055 -- the old +# ceiling silently clipped 598 of 672 values (89%). See +# docs/kwyk_mapping_verification.md, discrepancy D3. +_CONCRETE_P_MIN: float = 0.01 +_CONCRETE_P_MAX: float = 0.99 + +# F.normalize uses eps=1e-12 internally; reproduce it in the numpy parity path. +_NORMALIZE_EPS: float = 1e-12 + +# Default parity tolerances (overridable via CLI). +_DEFAULT_ATOL: float = 1e-5 +_DEFAULT_RTOL: float = 1e-6 + +# TF variable name templates. ``{i}`` is the layer index (base auto-detected). +_TF_CONV_V = "layer_{i}/conv3d/v" +_TF_CONV_G = "layer_{i}/conv3d/g" +_TF_CONV_KERNEL_A = "layer_{i}/conv3d/kernel_a" +_TF_CONV_BIAS_M = "layer_{i}/conv3d/bias_m" +_TF_CONV_BIAS_A = "layer_{i}/conv3d/bias_a" +_TF_CONCRETE_P = "layer_{i}/concrete_dropout/p" + +# The output layer lives in its OWN namespace, ``logits/``, and is itself a +# VWN conv with bias and no concrete_dropout sibling. Verified against the +# actual ``neuronets/kwyk:latest-cpu`` container (all_50_bvwn_multi_prior, +# timestamp 1556816070): ``logits/conv3d/v [1,1,1,96,50]``, ``g``, +# ``kernel_a``, ``bias_m``, ``bias_a`` -- 48 variables total including +# ``global_step``, matching ARCHITECTURE.md's counts exactly. All five map +# 1:1 onto the FFG ``classifier`` (kwyk_meshnet.py), so nothing is discarded. +_TF_LOGITS_V = "logits/conv3d/v" +_TF_LOGITS_G = "logits/conv3d/g" +_TF_LOGITS_KERNEL_A = "logits/conv3d/kernel_a" +_TF_LOGITS_BIAS_M = "logits/conv3d/bias_m" +_TF_LOGITS_BIAS_A = "logits/conv3d/bias_a" + + +class ConversionError(RuntimeError): + """Raised when the TF->PyTorch mapping cannot be completed safely.""" + + +# --- TF variable loading ------------------------------------------------------- + + +def load_tf_variables(npz_path: Path) -> dict[str, np.ndarray]: + """Load TF variables from a pre-extracted ``.npz``. + + Keys must be exact TF variable names without the trailing ``:0`` (e.g. + ``layer_1/conv3d/v``). This is the preferred path; it needs no TensorFlow. + + Parameters + ---------- + npz_path : Path + Path to the ``.npz`` produced by the extraction recipe (kwyk issue #15). + + Returns + ------- + dict[str, np.ndarray] + Mapping from TF variable name to array (float32-cast on read). + """ + with np.load(npz_path) as data: + return {key: np.asarray(data[key], dtype=np.float32) for key in data.files} + + +def load_tf_variables_from_savedmodel(tf_path: Path) -> dict[str, np.ndarray]: + """Load TF variables directly from a SavedModel/checkpoint. + + Requires ``tensorflow``. If it is not importable, raise with a message + pointing the user to the ``--npz`` path instead. + """ + try: + import tensorflow as tf # noqa: PLC0415 (optional dependency by design) + except ImportError as exc: # pragma: no cover - environment-dependent + raise ConversionError( + "tensorflow is not installed; cannot read a SavedModel directly. " + "Extract variables to an .npz (see kwyk issue #15) and pass --npz." + ) from exc + + reader = tf.train.load_checkpoint(str(tf_path)) + shapes = reader.get_variable_to_shape_map() + out: dict[str, np.ndarray] = {} + for name in shapes: + clean = name[:-2] if name.endswith(":0") else name + out[clean] = np.asarray(reader.get_tensor(name), dtype=np.float32) + return out + + +# --- Index-base detection ------------------------------------------------------ + + +def detect_layer_index_base(tf_vars: dict[str, np.ndarray], n_layers: int) -> int: + """Detect whether TF variable names are 0-based or 1-based, and validate. + + The PyTorch model uses ``layer_0 .. layer_{n_layers-1}``. The TF SavedModel + documented in ARCHITECTURE.md uses ``layer_1 ..``. An off-by-one here loads + every layer's weights into the wrong layer without necessarily erroring, so + this is asserted explicitly rather than assumed. + + Returns + ------- + int + The detected base (0 or 1). + """ + layer_indices: set[int] = set() + pattern = re.compile(r"^layer_(\d+)/") + for key in tf_vars: + match = pattern.match(key) + if match: + layer_indices.add(int(match.group(1))) + + if not layer_indices: + raise ConversionError( + "No 'layer_/...' variables found in the TF checkpoint. " + "Check the .npz keys use exact TF names without ':0'." + ) + + lo, hi = min(layer_indices), max(layer_indices) + span = hi - lo + 1 + # The output layer is NOT in this namespace (it is ``logits/conv3d/...``, + # verified against the real container), so the layer_{i} span must equal + # the hidden-layer count exactly. + if span != n_layers: + raise ConversionError( + f"TF checkpoint spans {span} layers (indices {lo}..{hi}) but the " + f"PyTorch model has {n_layers}. Refusing to convert on a mismatch." + ) + if lo not in (0, 1): + raise ConversionError(f"Unexpected TF layer base index {lo}; expected 0 or 1.") + return lo + + +# --- Shape conversion helpers -------------------------------------------------- + + +def _conv_tf_to_torch(arr: np.ndarray) -> np.ndarray: + """TF conv filter ``[k,k,k,in,out]`` -> torch ``[out,in,k,k,k]``.""" + if arr.ndim != 5: + raise ConversionError( + f"Expected 5-D conv filter [k,k,k,in,out], got shape {arr.shape}." + ) + return np.ascontiguousarray(np.transpose(arr, _CONV_PERM), dtype=np.float32) + + +def _g_tf_to_torch(arr: np.ndarray, out_channels: int) -> np.ndarray: + """TF ``g`` (``[1,1,1,1,out]`` or ``[out]``) -> torch ``(out,1,1,1,1)``. + + ``np.transpose`` on a 1-D array is a no-op, so the two cases are branched + explicitly on ``ndim`` rather than transposed blindly. + """ + if arr.ndim == 1: + if arr.size != out_channels: + raise ConversionError( + f"g has {arr.size} elements but out_channels={out_channels}." + ) + # Every published kwyk checkpoint stores g as 5-D [1,1,1,1,out]; a + # 1-D g is unusual enough to flag rather than silently reshape + # (docs/kwyk_mapping_verification.md, discrepancy D4). + logger.info( + "g is 1-D (%d elements); real kwyk checkpoints store 5-D g -- " + "reshaping, but verify the source checkpoint.", + arr.size, + ) + g = arr.reshape(out_channels, 1, 1, 1, 1) + elif arr.ndim == 5: + g = np.transpose(arr, _CONV_PERM) + else: + raise ConversionError(f"Unexpected g shape {arr.shape}; expected [out] or 5-D.") + if g.shape != (out_channels, 1, 1, 1, 1): + raise ConversionError( + f"Converted g shape {g.shape} != expected {(out_channels, 1, 1, 1, 1)}." + ) + return np.ascontiguousarray(g, dtype=np.float32) + + +def _p_to_logit(p: np.ndarray, name: str = "concrete_dropout/p") -> np.ndarray: + """Concrete-dropout probability -> logit, clamped to the forward clamp range. + + Any value outside ``[_CONCRETE_P_MIN, _CONCRETE_P_MAX]`` is clipped and + the count is logged as a warning -- clipping is a one-signed distortion, + never silent (docs/kwyk_mapping_verification.md, discrepancy D3). With the + current [0.01, 0.99] range, no published kwyk value is clipped. + """ + p32 = p.astype(np.float32) + n_clipped = int(((p32 < _CONCRETE_P_MIN) | (p32 > _CONCRETE_P_MAX)).sum()) + if n_clipped: + logger.warning( + "%s: %d of %d values fall outside [%.2f, %.2f] and were clipped " + "(source range [%.6f, %.6f]) -- the converted dropout rates are " + "distorted at these channels.", + name, + n_clipped, + p32.size, + _CONCRETE_P_MIN, + _CONCRETE_P_MAX, + float(p32.min()), + float(p32.max()), + ) + p_clamped = np.clip(p32, _CONCRETE_P_MIN, _CONCRETE_P_MAX) + return np.log(p_clamped / (1.0 - p_clamped)).astype(np.float32) + + +# --- State-dict construction --------------------------------------------------- + + +def build_state_dict( + model: torch.nn.Module, + tf_vars: dict[str, np.ndarray], + *, + create_bias: bool, +) -> dict[str, torch.Tensor]: + """Build a PyTorch state_dict from TF variables, matching model parameters. + + Parameters + ---------- + model : torch.nn.Module + A freshly constructed ``KWYKMeshNet`` (defines target shapes and the + layer count via ``_n_layers``). + tf_vars : dict[str, np.ndarray] + TF variables keyed by exact name (no ``:0``). + create_bias : bool + If True (the default in :func:`convert` -- every published kwyk + checkpoint carries biases for every conv layer), emit + ``conv.bias_m`` / ``conv.bias_a`` entries from the TF biases. The + model must have been built with bias-capable conv layers + (``bias=True``); registering the parameters is sufficient because + FFGConv3d.forward already consumes bias_m/bias_a. Setting this False + (CLI: ``--drop-bias``) discards data measured to change the output + logits by max ~20 on the published MAP weights. + + Returns + ------- + dict[str, torch.Tensor] + A state_dict intended for ``load_state_dict(..., strict=True)``. + """ + n_layers = int(getattr(model, "_n_layers")) + base = detect_layer_index_base(tf_vars, n_layers) + state: dict[str, torch.Tensor] = {} + + for torch_i in range(n_layers): + tf_i = torch_i + base + conv_prefix = f"layer_{torch_i}.conv" + + v = tf_vars[_TF_CONV_V.format(i=tf_i)] + g = tf_vars[_TF_CONV_G.format(i=tf_i)] + kernel_a = tf_vars[_TF_CONV_KERNEL_A.format(i=tf_i)] + + v_t = _conv_tf_to_torch(v) + out_channels = v_t.shape[0] + state[f"{conv_prefix}.v"] = torch.from_numpy(v_t) + state[f"{conv_prefix}.g"] = torch.from_numpy(_g_tf_to_torch(g, out_channels)) + state[f"{conv_prefix}.kernel_a"] = torch.from_numpy(_conv_tf_to_torch(kernel_a)) + + bias_m_key = _TF_CONV_BIAS_M.format(i=tf_i) + bias_a_key = _TF_CONV_BIAS_A.format(i=tf_i) + has_bias = bias_m_key in tf_vars and bias_a_key in tf_vars + if create_bias: + if not has_bias: + raise ConversionError( + f"Bias import requested but layer {tf_i} lacks " + "bias_m/bias_a in the TF checkpoint." + ) + state[f"{conv_prefix}.bias_m"] = torch.from_numpy( + np.ascontiguousarray(tf_vars[bias_m_key], dtype=np.float32) + ) + state[f"{conv_prefix}.bias_a"] = torch.from_numpy( + np.ascontiguousarray(tf_vars[bias_a_key], dtype=np.float32) + ) + elif has_bias: + logger.warning( + "Layer %d has TF bias_m/bias_a but bias import is disabled " + "(--drop-bias): DROPPING them. Measured on the published MAP " + "weights this shifts output logits by max ~20 (vs ~1e-4 with " + "biases kept) -- the result is NOT faithful to the original.", + tf_i, + ) + + # Confirmed: ConcreteDropout3d stores its learnable logit as + # ``p_logit`` (vwn_layers.py:197), so ``layer_{i}.dropout.p_logit`` is + # the correct key -- but only for a model built with + # dropout_type="concrete". Against the default "bernoulli" variant + # (parameter-free nn.Dropout3d) the strict load rejects it, which is + # the intended loud failure rather than a silent mismatch. + p_key = _TF_CONCRETE_P.format(i=tf_i) + if p_key in tf_vars: + state[f"layer_{torch_i}.dropout.p_logit"] = torch.from_numpy( + _p_to_logit(tf_vars[p_key], name=p_key) + ) + + _add_classifier(state, tf_vars) + return state + + +def _add_classifier( + state: dict[str, torch.Tensor], + tf_vars: dict[str, np.ndarray], +) -> None: + """Map the TF VWN ``logits/`` layer 1:1 onto the FFG ``classifier``. + + The TF output layer (``logits/conv3d/...``, its own namespace -- verified + against the real container) is a full VWN conv, and + ``KWYKMeshNet.classifier`` is an ``FFGConv3d`` as well, so all five + variables map directly with **no information loss** -- including the + ``kernel_a``/``bias_a`` sigmas that a plain-conv classifier could not + represent and whose absence made MC uncertainty under-dispersed at the + output (docs/kwyk_mapping_verification.md, discrepancy D2). + """ + required = ( + _TF_LOGITS_V, + _TF_LOGITS_G, + _TF_LOGITS_KERNEL_A, + _TF_LOGITS_BIAS_M, + _TF_LOGITS_BIAS_A, + ) + missing = [k for k in required if k not in tf_vars] + if missing: + raise ConversionError( + f"Checkpoint has no logits layer (missing {missing}). The model's " + "classifier cannot be filled, and a partial state_dict would fail " + "the strict load anyway." + ) + + v_t = _conv_tf_to_torch(tf_vars[_TF_LOGITS_V]) + out_channels = v_t.shape[0] + state["classifier.v"] = torch.from_numpy(v_t) + state["classifier.g"] = torch.from_numpy( + _g_tf_to_torch(tf_vars[_TF_LOGITS_G], out_channels) + ) + state["classifier.kernel_a"] = torch.from_numpy( + _conv_tf_to_torch(tf_vars[_TF_LOGITS_KERNEL_A]) + ) + state["classifier.bias_m"] = torch.from_numpy( + np.ascontiguousarray(tf_vars[_TF_LOGITS_BIAS_M], dtype=np.float32) + ) + state["classifier.bias_a"] = torch.from_numpy( + np.ascontiguousarray(tf_vars[_TF_LOGITS_BIAS_A], dtype=np.float32) + ) + + +# --- Parity check -------------------------------------------------------------- + + +def _kernel_m_numpy(v_torch: np.ndarray, g_torch: np.ndarray) -> np.ndarray: + """Reproduce ``FFGConv3d.kernel_m`` = ``g * normalize(v.flatten(1))`` in numpy. + + Mirrors ``F.normalize(self.v.flatten(1), dim=1).view_as(self.v)`` with the + same L2 norm over dims (in, k, k, k) per output channel and eps=1e-12. + """ + out = v_torch.shape[0] + v_flat = v_torch.reshape(out, -1) + norms = np.linalg.norm(v_flat, axis=1, keepdims=True) + norms = np.maximum(norms, _NORMALIZE_EPS) + v_norm = (v_flat / norms).reshape(v_torch.shape) + return (g_torch * v_norm).astype(np.float32) + + +def offline_parity_check( + model: torch.nn.Module, + tf_vars: dict[str, np.ndarray], + *, + atol: float, + rtol: float, + seed: int = 0, +) -> None: + """Verify parameter mapping without a TF runtime, layer by layer. + + For each conv layer, recompute ``kernel_m`` from the TF ``v``/``g`` in numpy, + run a deterministic ``F.conv3d`` mean path, and compare against the model's + own ``kernel_m`` conv on the same random input. This isolates the mapping + (axis order, normalization, bias) from every other confound. + + Scope: this validates the mapping *transform* (axis permutation, weight-norm, + bias arithmetic), not source-data integrity. Corruption identical on both the + recompute and the load path would move together and pass. For end-to-end + validation against the original model outputs, run the full-graph comparison + in ``scripts/kwyk_reproduction/05_compare_kwyk.py`` (Stage 3). + + Raises + ------ + ConversionError + If any layer's mean-path output diverges beyond tolerance. + """ + torch.manual_seed(seed) + n_layers = int(getattr(model, "_n_layers")) + base = detect_layer_index_base(tf_vars, n_layers) + model.eval() + + # Hidden convs plus the FFG classifier -- the output layer is part of the + # mapping and must be part of the check. + entries: list[tuple[str, torch.nn.Module, str, str]] = [ + ( + f"layer_{torch_i} (TF layer_{torch_i + base})", + getattr(model, f"layer_{torch_i}").conv, + _TF_CONV_V.format(i=torch_i + base), + _TF_CONV_G.format(i=torch_i + base), + ) + for torch_i in range(n_layers) + ] + entries.append( + ("classifier (TF logits)", model.classifier, _TF_LOGITS_V, _TF_LOGITS_G) + ) + + with torch.no_grad(): + for name, conv, v_key, g_key in entries: + v_t = _conv_tf_to_torch(tf_vars[v_key]) + g_t = _g_tf_to_torch(tf_vars[g_key], v_t.shape[0]) + kernel_m_ref = torch.from_numpy(_kernel_m_numpy(v_t, g_t)) + + in_ch = v_t.shape[1] + x = torch.randn(1, in_ch, 8, 8, 8) + + bias_m = conv.bias_m if conv.bias_m is not None else None + out_ref = F.conv3d( + x, kernel_m_ref, bias_m, conv.stride, conv.padding, conv.dilation + ) + out_model = F.conv3d( + x, conv.kernel_m, bias_m, conv.stride, conv.padding, conv.dilation + ) + + if not torch.allclose(out_ref, out_model, atol=atol, rtol=rtol): + max_abs = (out_ref - out_model).abs().max().item() + raise ConversionError( + f"Parity failed at {name}: " + f"max_abs_diff={max_abs:.3e} exceeds atol={atol:.1e}. " + "Likely an axis-order, normalization, or bias mismatch." + ) + logger.info( + "Offline parity check passed for all %d conv layers (incl. classifier).", + len(entries), + ) + + +# --- Orchestration ------------------------------------------------------------- + + +def convert( + model: torch.nn.Module, + *, + npz: Path | None = None, + tf_path: Path | None = None, + create_bias: bool = True, + run_parity: bool = True, + atol: float = _DEFAULT_ATOL, + rtol: float = _DEFAULT_RTOL, +) -> torch.nn.Module: + """Load TF variables, build a strict state_dict, verify parity, and load it. + + Exactly one of ``npz`` or ``tf_path`` must be given. + + ``create_bias`` defaults to True: every published kwyk checkpoint carries + ``bias_m``/``bias_a`` for every conv layer, and dropping them changes the + output logits by max ~20 (docs/kwyk_mapping_verification.md, D1). Pass + False only deliberately, with a model built ``bias=False``. + """ + if (npz is None) == (tf_path is None): + raise ConversionError("Provide exactly one of --npz or --tf-path.") + + tf_vars = ( + load_tf_variables(npz) if npz else load_tf_variables_from_savedmodel(tf_path) + ) + + state = build_state_dict(model, tf_vars, create_bias=create_bias) + # strict=True is a free correctness gate: any missing/unexpected key + # (a name typo, a dropped bias, a wrong layer count) fails loudly here. + model.load_state_dict(state, strict=True) + logger.info("Loaded converted weights into model (strict=True).") + + # Parity runs AFTER loading: it independently recomputes kernel_m from the + # raw TF v/g and checks the loaded model reproduces it. Running before the + # load would compare against random init and always fail. + if run_parity: + offline_parity_check(model, tf_vars, atol=atol, rtol=rtol) + + return model + + +def _build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + src = parser.add_mutually_exclusive_group(required=True) + src.add_argument("--npz", type=Path, help="Pre-extracted TF variables (preferred).") + src.add_argument( + "--tf-path", type=Path, help="SavedModel/checkpoint (requires tensorflow)." + ) + parser.add_argument( + "--drop-bias", + action="store_true", + help=( + "Discard the TF conv biases and build a bias-free model. Every " + "published kwyk checkpoint carries bias_m/bias_a for every conv " + "layer, and dropping them shifts output logits by max ~20 on the " + "MAP weights -- so the DEFAULT is to keep them (model built with " + "bias=True). This flag is an explicit, logged opt-out." + ), + ) + parser.add_argument( + "--dropout-type", + default="bernoulli", + choices=("bernoulli", "concrete"), + help=( + "Must be 'concrete' to import a checkpoint containing " + "concrete_dropout/p: only that variant has a learnable " + "dropout.p_logit parameter (the bernoulli variant uses a " + "parameter-free nn.Dropout3d, so the key would be unexpected)." + ), + ) + parser.add_argument("--n-classes", type=int, required=True) + parser.add_argument( + "--filters", + type=int, + default=96, + help="Hidden-layer filter count. The published kwyk models use 96.", + ) + parser.add_argument( + "--receptive-field", + type=int, + default=37, + choices=(37, 67, 129), + help="Dilation schedule selector. The published kwyk models use 37.", + ) + parser.add_argument("--out", type=Path, required=True, help="Output .pth path.") + parser.add_argument("--no-parity", action="store_true", help="Skip parity check.") + parser.add_argument("--atol", type=float, default=_DEFAULT_ATOL) + parser.add_argument("--rtol", type=float, default=_DEFAULT_RTOL) + return parser + + +def _sha256_of_file(path: Path) -> str | None: + """SHA-256 hex digest of a regular file; None for a directory/missing path.""" + if not path.is_file(): + return None + h = hashlib.sha256() + with open(path, "rb") as fh: + for chunk in iter(lambda: fh.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + +def _write_provenance(out: Path, args: argparse.Namespace, create_bias: bool) -> Path: + """Write a ``.provenance.json`` sidecar next to the weights. + + The MAP and BD checkpoints are structurally indistinguishable (identical + key sets and shapes -- docs/kwyk_mapping_verification.md, D5), so the + source path and its content hash are the only durable record of which + model a ``.pth`` came from. The architecture args are recorded because + the state_dict does not store them (dropout_type, receptive_field, etc. + are constructor-time-only). + """ + source = args.npz or args.tf_path + provenance = { + "source": str(source), + "source_sha256": _sha256_of_file(source), + "model_registry_name": "kwyk_meshnet", + "n_classes": args.n_classes, + "filters": args.filters, + "receptive_field": args.receptive_field, + "dropout_type": args.dropout_type, + "bias": create_bias, + "converter": "nobrainer.datasets.convert_kwyk", + "torch_version": torch.__version__, + "numpy_version": np.__version__, + "date": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + prov_path = out.with_name(out.stem + ".provenance.json") + prov_path.write_text(json.dumps(provenance, indent=2) + "\n") + return prov_path + + +def main(argv: list[str] | None = None) -> int: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + args = _build_arg_parser().parse_args(argv) + create_bias = not args.drop_bias + + # Imported here so the module's helpers are usable without nobrainer present + # (e.g. in unit tests that construct a stub model). + from nobrainer.models import get as get_model # noqa: PLC0415 + + # The registry name is "kwyk_meshnet" (nobrainer/models/__init__.py); there + # is no "kwyk" export. bias must mirror create_bias, because a strict + # state_dict load rejects conv.bias_m/bias_a against a bias-free model + # (and vice versa reports them missing). + model = get_model("kwyk_meshnet")( + n_classes=args.n_classes, + filters=args.filters, + receptive_field=args.receptive_field, + dropout_type=args.dropout_type, + bias=create_bias, + ) + convert( + model, + npz=args.npz, + tf_path=args.tf_path, + create_bias=create_bias, + run_parity=not args.no_parity, + atol=args.atol, + rtol=args.rtol, + ) + torch.save(model.state_dict(), args.out) + logger.info("Wrote converted weights to %s", args.out) + prov_path = _write_provenance(args.out, args, create_bias) + logger.info("Wrote provenance sidecar to %s", prov_path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/nobrainer/models/bayesian/kwyk_meshnet.py b/nobrainer/models/bayesian/kwyk_meshnet.py index a0fd6319..f50319e6 100644 --- a/nobrainer/models/bayesian/kwyk_meshnet.py +++ b/nobrainer/models/bayesian/kwyk_meshnet.py @@ -38,6 +38,7 @@ def __init__( dilation: int, dropout_rate: float, sigma_init: float, + bias: bool = False, ) -> None: super().__init__() self.conv = FFGConv3d( @@ -46,7 +47,7 @@ def __init__( kernel_size=3, padding=dilation, dilation=dilation, - bias=False, + bias=bias, sigma_init=sigma_init, ) self.dropout = nn.Dropout3d(p=dropout_rate) @@ -75,6 +76,7 @@ def __init__( sigma_init: float, concrete_temperature: float = 0.02, concrete_init_p: float = 0.9, + bias: bool = False, ) -> None: super().__init__() self.conv = FFGConv3d( @@ -83,7 +85,7 @@ def __init__( kernel_size=3, padding=dilation, dilation=dilation, - bias=False, + bias=bias, sigma_init=sigma_init, ) self.dropout = ConcreteDropout3d( @@ -131,6 +133,18 @@ class KWYKMeshNet(nn.Module): Temperature for concrete dropout (default 0.02). concrete_init_p : float Initial dropout probability for concrete dropout (default 0.9). + bias : bool + Whether the hidden-layer FFG convolutions carry a bias term + (``bias_m`` / ``bias_a``). Default ``False`` (the from-scratch + training default). Note the **published kwyk checkpoints DO carry + ``bias_m``/``bias_a`` for every conv layer** — verified against the + ``neuronets/kwyk`` container, see + ``docs/kwyk_mapping_verification.md`` — so importing them requires + ``bias=True`` (``nobrainer/datasets/convert_kwyk.py`` builds the + model that way by default); a strict state_dict load rejects the + bias keys against a bias-free model. The output ``classifier`` is + unaffected by this flag: it always carries its bias, matching both + the TF ``logits/`` layer and the previous ``nn.Conv3d`` behaviour. """ def __init__( @@ -144,6 +158,7 @@ def __init__( sigma_init: float = 1e-4, concrete_temperature: float = 0.02, concrete_init_p: float = 0.9, + bias: bool = False, ) -> None: super().__init__() if receptive_field not in _DILATION_SCHEDULES: @@ -165,6 +180,7 @@ def __init__( sigma_init, concrete_temperature, concrete_init_p, + bias=bias, ) else: layer = _VWNLayerBernoulli( @@ -173,10 +189,21 @@ def __init__( dil, dropout_rate, sigma_init, + bias=bias, ) setattr(self, f"layer_{i}", layer) - self.classifier = nn.Conv3d(filters, n_classes, kernel_size=1) + # The output layer is itself a VWN (FFG) conv, exactly like the TF + # original's ``logits/conv3d/*`` layer (a full VWN conv with bias -- + # verified against the neuronets/kwyk container, see + # docs/kwyk_mapping_verification.md, discrepancy D2). A plain + # nn.Conv3d here could represent only the mean path, making MC + # uncertainty systematically under-dispersed at the output. bias is + # unconditional: the TF logits layer always has one, and so did the + # previous nn.Conv3d default. + self.classifier = FFGConv3d( + filters, n_classes, kernel_size=1, bias=True, sigma_init=sigma_init + ) def forward( self, @@ -213,7 +240,9 @@ def forward( h = x for i in range(self._n_layers): h = getattr(self, f"layer_{i}")(h, mc_vwn=mc_vwn, mc_dropout=mc_dropout) - return self.classifier(h) + # The classifier follows the same VWN sampling switch as the hidden + # convs: stochastic under mc_vwn=True, mean path under mc_vwn=False. + return self.classifier(h, mc=mc_vwn) def kl_divergence(self) -> torch.Tensor: """Sum KL divergence from all VWN conv layers.""" @@ -242,9 +271,17 @@ def kwyk_meshnet( sigma_init: float = 1e-4, concrete_temperature: float = 0.02, concrete_init_p: float = 0.9, + bias: bool = False, **kwargs, ) -> KWYKMeshNet: - """Factory function for :class:`KWYKMeshNet`.""" + """Factory function for :class:`KWYKMeshNet`. + + ``bias`` defaults to ``False`` (the from-scratch training default). The + published kwyk checkpoints DO carry ``bias_m``/``bias_a`` for every conv + layer (verified against the ``neuronets/kwyk`` container -- see + ``docs/kwyk_mapping_verification.md``), so pass ``bias=True`` when + importing them; ``nobrainer/datasets/convert_kwyk.py`` does so by default. + """ return KWYKMeshNet( n_classes=n_classes, in_channels=in_channels, @@ -255,6 +292,7 @@ def kwyk_meshnet( sigma_init=sigma_init, concrete_temperature=concrete_temperature, concrete_init_p=concrete_init_p, + bias=bias, ) diff --git a/nobrainer/models/bayesian/vwn_layers.py b/nobrainer/models/bayesian/vwn_layers.py index b823e8f1..f81d5083 100644 --- a/nobrainer/models/bayesian/vwn_layers.py +++ b/nobrainer/models/bayesian/vwn_layers.py @@ -30,6 +30,16 @@ import torch.nn as nn import torch.nn.functional as F +# Clamp range for concrete-dropout probabilities, applied on every read of +# ``ConcreteDropout3d.p``. Widened from the original [0.05, 0.95]: the +# published kwyk SSD checkpoint stores trained p values up to 0.954055, so a +# 0.95 ceiling silently clipped 89% of them (598/672) on weight import +# (docs/kwyk_mapping_verification.md, discrepancy D3). The converter in +# ``nobrainer/datasets/convert_kwyk.py`` mirrors these bounds; a unit test +# asserts the two stay in sync. +CONCRETE_P_MIN: float = 0.01 +CONCRETE_P_MAX: float = 0.99 + class FFGConv3d(nn.Module): """3-D convolution with Variational Weight Normalization + learned sigma. @@ -198,8 +208,13 @@ def __init__( @property def p(self) -> torch.Tensor: - """Per-filter dropout probabilities, clamped to [0.05, 0.95].""" - return torch.sigmoid(self.p_logit).clamp(0.05, 0.95) + """Per-filter dropout probabilities, clamped to the module range. + + The clamp bounds are the module constants ``CONCRETE_P_MIN`` / + ``CONCRETE_P_MAX``; see their definition for why the range must be + wide enough to represent the published kwyk checkpoints. + """ + return torch.sigmoid(self.p_logit).clamp(CONCRETE_P_MIN, CONCRETE_P_MAX) def forward(self, x: torch.Tensor, mc: bool = True) -> torch.Tensor: """Apply concrete dropout (Eq. 10). diff --git a/nobrainer/tests/unit/test_convert_kwyk.py b/nobrainer/tests/unit/test_convert_kwyk.py new file mode 100644 index 00000000..a149ea31 --- /dev/null +++ b/nobrainer/tests/unit/test_convert_kwyk.py @@ -0,0 +1,381 @@ +"""Unit tests for nobrainer.datasets.convert_kwyk. + +No TensorFlow or network dependency: a stub model matching the KWYKMeshNet +interface is used as ground truth. Fabricated TF-layout variables are derived +from it, converted back, and checked for round-trip equality. + +If these tests fail after wiring the real model, the likely causes are: +- ConcreteDropout logit attribute is not `dropout.p_logit` (update the template). +- The kwyk factory cannot build bias-capable convs (see test_create_bias_*). +""" + +from __future__ import annotations + +import math +from pathlib import Path + +import numpy as np +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from nobrainer.datasets import convert_kwyk as ck + +# --- Stub model mirroring the real KWYKMeshNet interface ----------------------- + + +class _StubFFGConv3d(nn.Module): + def __init__( + self, in_ch: int, out_ch: int, k: int = 3, dilation: int = 1, bias: bool = False + ) -> None: + super().__init__() + self.in_channels, self.out_channels = in_ch, out_ch + self.kernel_size, self.stride = k, 1 + self.padding, self.dilation = dilation, dilation + ws = (out_ch, in_ch, k, k, k) + self.v = nn.Parameter(torch.empty(ws)) + nn.init.kaiming_normal_(self.v) + self.g = nn.Parameter(torch.full((out_ch, 1, 1, 1, 1), math.sqrt(2.0))) + self.kernel_a = nn.Parameter(torch.full(ws, 1e-4)) + if bias: + self.bias_m = nn.Parameter(torch.zeros(out_ch)) + self.bias_a = nn.Parameter(torch.full((out_ch,), 1e-4)) + else: + self.register_parameter("bias_m", None) + self.register_parameter("bias_a", None) + + @property + def kernel_m(self) -> torch.Tensor: + v_norm = F.normalize(self.v.flatten(1), dim=1).view_as(self.v) + return self.g * v_norm + + +class _StubDropout(nn.Module): + def __init__(self, out_ch: int) -> None: + super().__init__() + self.p_logit = nn.Parameter(torch.zeros(out_ch)) + + +class _StubLayer(nn.Module): + def __init__(self, in_ch: int, out_ch: int, dilation: int, bias: bool) -> None: + super().__init__() + self.conv = _StubFFGConv3d(in_ch, out_ch, 3, dilation, bias) + self.dropout = _StubDropout(out_ch) + + +class _StubKWYK(nn.Module): + def __init__(self, bias: bool = False) -> None: + super().__init__() + dilations = [1, 2, 4] + self._n_layers = len(dilations) + for i, d in enumerate(dilations): + in_ch = 1 if i == 0 else 8 + setattr(self, f"layer_{i}", _StubLayer(in_ch, 8, d, bias)) + # Mirror the real KWYKMeshNet: the classifier is itself an FFG conv + # (v/g/kernel_a/bias_m/bias_a), matching the TF logits/ layer 1:1. + # Without this the stub diverges from the real model in exactly the + # spot the converter must fill, and the suite would stay green while + # main() fails on the real architecture. + self.classifier = _StubFFGConv3d(8, 2, k=1, dilation=1, bias=True) + + +# --- Fixtures ------------------------------------------------------------------ + + +@pytest.fixture +def ground_truth() -> _StubKWYK: + torch.manual_seed(1) + return _StubKWYK(bias=True) + + +@pytest.fixture +def tf_vars(ground_truth: _StubKWYK) -> dict[str, np.ndarray]: + """Fabricate 1-based TF-layout variables from the ground-truth model.""" + out: dict[str, np.ndarray] = {} + for i in range(ground_truth._n_layers): + conv = getattr(ground_truth, f"layer_{i}").conv + # torch [out,in,k,k,k] -> TF [k,k,k,in,out] (inverse of _CONV_PERM) + out[f"layer_{i + 1}/conv3d/v"] = np.transpose( + conv.v.detach().numpy(), (2, 3, 4, 1, 0) + ) + out[f"layer_{i + 1}/conv3d/g"] = np.transpose( + conv.g.detach().numpy(), (2, 3, 4, 1, 0) + ) + out[f"layer_{i + 1}/conv3d/kernel_a"] = np.transpose( + conv.kernel_a.detach().numpy(), (2, 3, 4, 1, 0) + ) + out[f"layer_{i + 1}/conv3d/bias_m"] = conv.bias_m.detach().numpy() + out[f"layer_{i + 1}/conv3d/bias_a"] = conv.bias_a.detach().numpy() + p = torch.sigmoid(getattr(ground_truth, f"layer_{i}").dropout.p_logit) + out[f"layer_{i + 1}/concrete_dropout/p"] = p.detach().numpy() + + # Output layer: its own ``logits/`` namespace, itself a full VWN conv + # with NO concrete_dropout sibling (verified against + # neuronets/kwyk:latest-cpu). All five variables map 1:1 onto the FFG + # classifier, so fabricate them from the stub classifier's own params. + clf = ground_truth.classifier + out["logits/conv3d/v"] = np.transpose(clf.v.detach().numpy(), (2, 3, 4, 1, 0)) + out["logits/conv3d/g"] = np.transpose(clf.g.detach().numpy(), (2, 3, 4, 1, 0)) + out["logits/conv3d/kernel_a"] = np.transpose( + clf.kernel_a.detach().numpy(), (2, 3, 4, 1, 0) + ) + out["logits/conv3d/bias_m"] = clf.bias_m.detach().numpy() + out["logits/conv3d/bias_a"] = clf.bias_a.detach().numpy() + return out + + +@pytest.fixture +def npz_path(tf_vars: dict[str, np.ndarray], tmp_path: Path) -> Path: + path = tmp_path / "fake_tf.npz" + np.savez(path, **tf_vars) + return path + + +# --- Tests: happy path --------------------------------------------------------- + + +def test_convert_roundtrip_matches_ground_truth(ground_truth, npz_path): + """Converted model reproduces the ground-truth mean path within tolerance. + + Deliberately relies on convert()'s DEFAULT create_bias -- which must keep + biases, since every published checkpoint has them (report D1). + """ + dst = _StubKWYK(bias=True) + ck.convert(dst, npz=npz_path, run_parity=True) + x = torch.randn(1, 1, 8, 8, 8) + with torch.no_grad(): + a = ground_truth.layer_0.conv + b = dst.layer_0.conv + out_a = F.conv3d(x, a.kernel_m, a.bias_m, a.stride, a.padding, a.dilation) + out_b = F.conv3d(x, b.kernel_m, b.bias_m, b.stride, b.padding, b.dilation) + torch.testing.assert_close(out_a, out_b, atol=1e-6, rtol=1e-6) + + +def test_p_logit_recovered(ground_truth, npz_path): + dst = _StubKWYK(bias=True) + ck.convert(dst, npz=npz_path, create_bias=True, run_parity=False) + torch.testing.assert_close( + dst.layer_0.dropout.p_logit, + ground_truth.layer_0.dropout.p_logit, + atol=1e-5, + rtol=1e-5, + ) + + +def test_classifier_recovered(ground_truth, npz_path): + """The VWN logits layer maps 1:1 onto the FFG classifier -- all five params.""" + dst = _StubKWYK(bias=True) + ck.convert(dst, npz=npz_path, create_bias=True, run_parity=False) + for attr in ("v", "g", "kernel_a", "bias_m", "bias_a"): + torch.testing.assert_close( + getattr(dst.classifier, attr), + getattr(ground_truth.classifier, attr), + atol=1e-6, + rtol=1e-6, + ) + + +def test_missing_logits_layer_rejected(tf_vars, tmp_path): + """A checkpoint without the logits layer must fail loudly, not partially load.""" + hidden_only = {k: v for k, v in tf_vars.items() if not k.startswith("logits/")} + path = tmp_path / "no_logits.npz" + np.savez(path, **hidden_only) + dst = _StubKWYK(bias=True) + with pytest.raises(ck.ConversionError, match="no logits layer"): + ck.convert(dst, npz=path, create_bias=True, run_parity=False) + + +def test_detect_index_base_one(tf_vars): + assert ck.detect_layer_index_base(tf_vars, 3) == 1 + + +def test_g_1d_handled(): + g = np.random.randn(8).astype(np.float32) + out = ck._g_tf_to_torch(g, 8) + assert out.shape == (8, 1, 1, 1, 1) + + +def test_g_5d_handled(): + g = np.random.randn(1, 1, 1, 1, 8).astype(np.float32) + out = ck._g_tf_to_torch(g, 8) + assert out.shape == (8, 1, 1, 1, 1) + + +# --- Tests: guard rails (must fail loudly) ------------------------------------- + + +def test_create_bias_on_biasless_model_rejected(npz_path): + """Model built bias=False + bias import (the default) must fail strict load. + + The mismatch guard: importing biases into a bias-free model produces + unexpected keys, which strict=True rejects loudly. + """ + dst = _StubKWYK(bias=False) + with pytest.raises(RuntimeError): # strict load raises on unexpected keys + ck.convert(dst, npz=npz_path, create_bias=True, run_parity=False) + + +def test_layer_count_mismatch_caught(tf_vars): + with pytest.raises(ck.ConversionError): + ck.detect_layer_index_base(tf_vars, 5) # npz has 3 layers + + +def test_conv_wrong_ndim_rejected(): + with pytest.raises(ck.ConversionError): + ck._conv_tf_to_torch(np.zeros((3, 3, 3))) # 3-D, not 5-D + + +def test_parity_detects_transform_error(ground_truth, npz_path, monkeypatch): + """If the conv transpose is wrong, parity must catch it. + + The offline check validates the mapping transform (axis order, weight-norm), + not source-data integrity: corrupting a TF var identically on both the + recompute and load paths would move together and pass. So we inject a wrong + permutation and confirm parity fails loudly. + """ + dst = _StubKWYK(bias=True) + # Patch the conv permutation to an incorrect axis order. + monkeypatch.setattr(ck, "_CONV_PERM", (0, 1, 2, 3, 4)) + with pytest.raises(ck.ConversionError): + ck.convert(dst, npz=npz_path, create_bias=True, run_parity=True) + + +def test_both_sources_rejected(npz_path): + dst = _StubKWYK(bias=True) + with pytest.raises(ck.ConversionError, match="exactly one"): + ck.convert(dst, npz=npz_path, tf_path=npz_path, create_bias=True) + + +def test_no_source_rejected(): + dst = _StubKWYK(bias=True) + with pytest.raises(ck.ConversionError, match="exactly one"): + ck.convert(dst) + + +# --- Tests: discrepancy fixes (docs/kwyk_mapping_verification.md) -------------- + + +def test_clamp_constants_match_model(): + """The converter's p clamp MUST equal ConcreteDropout3d's (report D3). + + The two are literal copies (the converter avoids importing nobrainer at + module scope); this test is the sync mechanism -- widening one without + the other silently reintroduces the clipping bug. + """ + from nobrainer.models.bayesian.vwn_layers import CONCRETE_P_MAX, CONCRETE_P_MIN + + assert ck._CONCRETE_P_MIN == CONCRETE_P_MIN + assert ck._CONCRETE_P_MAX == CONCRETE_P_MAX + + +def test_published_p_range_survives_conversion(): + """p values up to 0.954 (the real SSD max) must NOT be clipped (report D3).""" + p = np.array([0.607895, 0.9, 0.953, 0.954055], dtype=np.float32) + logit = ck._p_to_logit(p) + recovered = 1.0 / (1.0 + np.exp(-logit)) + np.testing.assert_allclose(recovered, p, rtol=1e-5, atol=1e-6) + + +def test_main_end_to_end_writes_weights_and_provenance(tmp_path): + """main() converts a real (tiny) KWYKMeshNet checkpoint and records + provenance -- the sidecar is the only durable record of which source + model a .pth came from, since MAP and BD checkpoints are structurally + indistinguishable (report D5). Also exercises report D1's new default: + no bias flag passed, biases kept.""" + import json + + from nobrainer.models import get as get_model + + torch.manual_seed(3) + src = get_model("kwyk_meshnet")( + n_classes=2, filters=8, receptive_field=37, dropout_type="concrete", bias=True + ) + tf_vars: dict = {} + for i in range(src._n_layers): + conv = getattr(src, f"layer_{i}").conv + tf_vars[f"layer_{i + 1}/conv3d/v"] = np.transpose( + conv.v.detach().numpy(), (2, 3, 4, 1, 0) + ) + tf_vars[f"layer_{i + 1}/conv3d/g"] = np.transpose( + conv.g.detach().numpy(), (2, 3, 4, 1, 0) + ) + tf_vars[f"layer_{i + 1}/conv3d/kernel_a"] = np.transpose( + conv.kernel_a.detach().numpy(), (2, 3, 4, 1, 0) + ) + tf_vars[f"layer_{i + 1}/conv3d/bias_m"] = conv.bias_m.detach().numpy() + tf_vars[f"layer_{i + 1}/conv3d/bias_a"] = conv.bias_a.detach().numpy() + p = torch.sigmoid(getattr(src, f"layer_{i}").dropout.p_logit) + tf_vars[f"layer_{i + 1}/concrete_dropout/p"] = p.detach().numpy() + clf = src.classifier + tf_vars["logits/conv3d/v"] = np.transpose(clf.v.detach().numpy(), (2, 3, 4, 1, 0)) + tf_vars["logits/conv3d/g"] = np.transpose(clf.g.detach().numpy(), (2, 3, 4, 1, 0)) + tf_vars["logits/conv3d/kernel_a"] = np.transpose( + clf.kernel_a.detach().numpy(), (2, 3, 4, 1, 0) + ) + tf_vars["logits/conv3d/bias_m"] = clf.bias_m.detach().numpy() + tf_vars["logits/conv3d/bias_a"] = clf.bias_a.detach().numpy() + + npz = tmp_path / "tiny_kwyk.npz" + np.savez(npz, **tf_vars) + out = tmp_path / "converted.pth" + + rc = ck.main( + [ + "--npz", + str(npz), + "--n-classes", + "2", + "--filters", + "8", + "--receptive-field", + "37", + "--dropout-type", + "concrete", + "--out", + str(out), + ] + ) + assert rc == 0 + assert out.exists() + + prov_path = tmp_path / "converted.provenance.json" + assert prov_path.exists() + prov = json.loads(prov_path.read_text()) + assert prov["source"] == str(npz) + assert prov["source_sha256"] == ck._sha256_of_file(npz) + assert prov["bias"] is True # D1: kept by default, no flag passed + assert prov["dropout_type"] == "concrete" + + # The written weights round-trip into a fresh model and reproduce the + # source's deterministic forward exactly (mean path). + dst = get_model("kwyk_meshnet")( + n_classes=2, filters=8, receptive_field=37, dropout_type="concrete", bias=True + ) + dst.load_state_dict(torch.load(out, weights_only=True), strict=True) + src.eval() + dst.eval() + x = torch.randn(1, 1, 16, 16, 16) + with torch.no_grad(): + torch.testing.assert_close( + src(x, mc=False), dst(x, mc=False), atol=1e-6, rtol=1e-6 + ) + + +def test_real_classifier_contributes_mc_variance(): + """The FFG classifier must sample under mc=True and be exactly + deterministic under mc=False -- the point of report D2's fix (the old + plain-conv classifier contributed zero output variance).""" + from nobrainer.models.bayesian.vwn_layers import FFGConv3d + + torch.manual_seed(5) + clf = FFGConv3d(8, 2, kernel_size=1, bias=True, sigma_init=0.1) + clf.eval() + x = torch.randn(1, 8, 4, 4, 4) + with torch.no_grad(): + a = clf(x, mc=True) + b = clf(x, mc=True) + det1 = clf(x, mc=False) + det2 = clf(x, mc=False) + assert not torch.allclose(a, b), "mc=True must sample" + torch.testing.assert_close(det1, det2, atol=0.0, rtol=0.0) diff --git a/scripts/kwyk_reproduction/ARCHITECTURE.md b/scripts/kwyk_reproduction/ARCHITECTURE.md index fbf86873..be18abf6 100644 --- a/scripts/kwyk_reproduction/ARCHITECTURE.md +++ b/scripts/kwyk_reproduction/ARCHITECTURE.md @@ -112,7 +112,9 @@ Two terms per filter: - Layer 5: dilation=4 - Layer 6: dilation=8 - Layer 7: dilation=1 -- Layer 8 (logits): 1×1×1, 50 filters, Softmax +- Layer 8 (logits): 1×1×1, 50 filters, Softmax — "Layer 8" is positional + only; in the SavedModel its variables live under their own `logits/conv3d/*` + namespace (not `layer_8/`), and it is itself a full VWN conv with bias Receptive field = 37 voxels.