200-epoch retrain: 9 variants with FP32 + int8 ONNX + int8 PESQ findings

README.md

@@ -9,6 +9,9 @@ tags:
   - butterfly
   - monarch
   - structured-matrices
+  - onnx
+  - quantization
+  - int8
 datasets:
   - JacobLinCool/VoiceBank-DEMAND-16k
 metrics:
@@ -25,41 +28,72 @@ dense, [Butterfly](https://arxiv.org/abs/1903.05895), and
 
 Trained on
 [VoiceBank-DEMAND-16k](https://huggingface.co/datasets/JacobLinCool/VoiceBank-DEMAND-16k)
-for 50 epochs at batch 256, n\_fft 512, on a single GTX 1080 Ti.
+for 200 epochs at batch 256, n\_fft 512, on a single GTX 1080 Ti.
+
+Each variant ships in three formats:
+- `g_best` — PyTorch checkpoint (training-time weights)
+- `g_best_fp32.onnx` — streaming-shape FP32 ONNX (frame-by-frame, opset 17)
+- `g_best.onnx` — static int8 ONNX (QDQ format, per-channel weight quant,
+  MinMax calibration on 200 VBD-train utterances)
 
 ## Results
 
-| run               | best PESQ | params | Δ baseline | compression |
-| ----------------- | --------: | -----: | ---------: | ----------: |
-| `wide_monarch`    | **2.833** | 2.36 M |    +0.080  |       1.2× |
-| `butterfly_fc`    |     2.809 | 1.99 M |    +0.056  |       1.4× |
-| `butterfly_2blocks` |   2.780 | **0.36 M** | +0.027 |   **7.7×** |
-| `monarch_fc`      |     2.777 | 2.14 M |    +0.024  |       1.3× |
-| `monarch_full`    |     2.764 | 0.70 M |    +0.011  |       4.0× |
-| `baseline`        |     2.753 | 2.78 M |       —    |       1.0× |
-| `butterfly_ortho` |     2.733 | 0.19 M |   −0.020   |        15× |
-| `monarch_8`       |     2.732 | 0.36 M |   −0.021   |       7.7× |
-| `butterfly_full`  |     2.673 | 0.19 M |   −0.080   |        15× |
+PESQ measured on the full VBD test split (824 utterances). RTF (real-time
+factor) is for the int8 ONNX session under onnxruntime CPU; lower is faster.
+
+| run                 | params  | FP32 PESQ |  int8 PESQ |  Δ (FP32→int8) | int8 RTF |
+| ------------------- | ------: | --------: | ---------: | -------------: | -------: |
+| `wide_monarch`      |  2.36 M | **2.864** |      2.842 |        +0.021  |    0.166 |
+| `baseline`          |  2.78 M |     2.845 |      2.833 |        +0.012  |    0.452 |
+| `monarch_8`         |  0.36 M |     2.832 |      2.826 |        +0.006  | **0.025** |
+| `monarch_full`      |  0.70 M |     2.827 |  **2.848** |        −0.021  |    0.039 |
+| `monarch_fc`        |  2.14 M |     2.805 |      2.789 |        +0.016  |    0.448 |
+| `butterfly_2blocks` |  0.36 M |     2.805 |      2.202 |        +0.602  |    0.441 |
+| `butterfly_fc`      |  1.99 M |     2.799 |      2.494 |        +0.306  |    0.522 |
+| `butterfly_ortho`   |  0.19 M |     2.780 |      2.577 |        +0.203  |    0.232 |
+| `butterfly_full`    |  0.19 M |     2.772 |      2.128 |        +0.644  |    0.230 |
+
+## Key findings
+
+- **Monarch variants are essentially loss-free under int8** (|Δ| ≤ 0.021
+  across the board). A single `Einsum` per FC plus per-channel weight
+  quantization is genuinely friendly to int8 calibration.
+- **`wide_monarch` is the best deployment target** for quality: highest FP32
+  PESQ (2.864) with near-zero int8 loss (2.842). For speed-constrained
+  deployment, **`monarch_full` and `monarch_8`** trade ~0.02 PESQ for an
+  RTF of 0.04 / 0.025 — over 10× faster than the dense baseline.
+- **Butterfly with random init degrades catastrophically under int8**
+  (Δ up to 0.64 PESQ on `butterfly_full`). Int8 deployment with butterfly
+  factorizations should use `init=ortho`: `butterfly_ortho` loses 0.20
+  PESQ to int8 versus 0.64 for `butterfly_full` (same architecture, same
+  training data, only the init differs).
+- **Longer training makes randn-init butterfly *worse* on int8.** The same
+  `butterfly_full` config saw its int8 gap grow from 0.36 PESQ at 50
+  epochs → 0.64 PESQ at 200 epochs as twiddle factors drifted further
+  from orthogonality. Ortho-init butterfly does not show this regression.
 
 ## Layout
 
-Each subdirectory is one run, containing the saved generator (`g_best`)
-and the exact `config.json` it was trained with.
+Each subdirectory is one run, containing the saved generator (`g_best`),
+the streaming FP32 ONNX, the static int8 ONNX, and the exact `config.json`
+the run was trained with.
 
 ```
-baseline/{g_best,config.json}
-monarch_fc/{g_best,config.json}
-butterfly_fc/{g_best,config.json}
-monarch_full/{g_best,config.json}
-butterfly_full/{g_best,config.json}
-monarch_8/{g_best,config.json}
-butterfly_ortho/{g_best,config.json}
-butterfly_2blocks/{g_best,config.json}
-wide_monarch/{g_best,config.json}
+baseline/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+monarch_fc/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+butterfly_fc/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+monarch_full/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+butterfly_full/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+monarch_8/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+butterfly_ortho/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+butterfly_2blocks/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
+wide_monarch/{g_best,g_best_fp32.onnx,g_best.onnx,config.json}
 ```
 
 ## Loading
 
+### PyTorch checkpoint
+
 Clone the repo first (model classes live there):
 
 ```bash
@@ -87,7 +121,30 @@ model = NSNet2(AttrDict(cfg)).cuda().eval()
 model.load_state_dict(ckpt["generator"])
 ```
 
-End-to-end inference example is in `inference.py` and `analyze_sweep.ipynb`
+### ONNX (FP32 or int8)
+
+The ONNX models are streaming-shape: a single frame `(B, n_freq)` plus the
+GRU state `(num_layers, B, hidden)` per session call, threaded across
+frames. The end-to-end pipeline (RMS-norm → STFT → frame loop → iSTFT)
+is in `inference_onnx.py` in the source repo.
+
+```python
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+REPO = "claroche1/sparse-nsnet2-checkpoints"
+RUN  = "wide_monarch"
+
+# FP32:
+fp32_path = hf_hub_download(REPO, f"{RUN}/g_best_fp32.onnx")
+fp32_sess = ort.InferenceSession(fp32_path, providers=["CPUExecutionProvider"])
+
+# int8 (deployment):
+int8_path = hf_hub_download(REPO, f"{RUN}/g_best.onnx")
+int8_sess = ort.InferenceSession(int8_path, providers=["CPUExecutionProvider"])
+```
+
+End-to-end inference example with PESQ measurement is in `inference_onnx.py`
 in the source repo.
 
 ## Citations

baseline/config.json

@@ -22,6 +22,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

baseline/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3bea76d887cb6216acaecb721e9c7a6e318b8f0fcc0986e00cf306d39310865d
+oid sha256:558f79c6037f7f80a36abde43bedd480f9a1b60fd8ef788cd9dd5733f89d29dc
 size 11138351

baseline/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66e5b8f4a540df133802c6ff3f3e91eacfc305dfd3b5b37cb31c13870bf70afd
+size 2910073

baseline/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81292578df182d2a24bee37f78df2011a95dabd8abf95ef35ec14461b060cabb
+size 11140400

butterfly_2blocks/config.json

@@ -26,6 +26,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

butterfly_2blocks/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a43a489acb1f9aee7bdc51648a84e816077fa485dbc62c05b80cf73554ab3e3f
+oid sha256:d097f594afa144b90660f00184b40215f91640a50ac43b93e65d155516b090f6
 size 1465234

butterfly_2blocks/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df45fd5ae9f1ea40f726f6f4150a722b45ff2b267f4acdd5dc20a62c9c802244
+size 821735

butterfly_2blocks/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:285afd772ee4eb456078a86988e39aeebe25c57de521319160417254de51d8af
+size 1625605

butterfly_fc/config.json

@@ -24,6 +24,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

butterfly_fc/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:954b58e3f9665f184d9539fbaea01c3f6be599c17bbd49605ae4dafd51f2b4eb
+oid sha256:92224e592913ef8b7ba5d7d930702f2c44460cdf47f8f05bc973fdb19d7b51d7
 size 7984815

butterfly_fc/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a4d9953d294358f791013b9ab759f7b284f7777c1f3f6d8a1a0623364eed3b9
+size 2198876

butterfly_fc/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f52b948e0f9bb7e38f96a477b0f1224cd1d5cb14daabc9d049d4c1a4cee86728
+size 8032362

butterfly_full/config.json

@@ -26,6 +26,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

butterfly_full/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:161fd66cbe4d109a80d7862386ee66a6850ae4ef189c852335c557bce0b98e78
+oid sha256:91510093d1d1f9cba2903b8a9229c421255f4feac7b8cb6935ed0c260a809cdc
 size 748434

butterfly_full/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bef3f4279db6a0b90ab43fe0f1de0e9263db5897cf00f05323caf76220fd492
+size 476461

butterfly_full/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5422e6e3d6f4c54334e2b156eca1b4a44fc556f7e388f478b34f3a1cb3b918fc
+size 865393

butterfly_ortho/config.json

@@ -27,6 +27,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

butterfly_ortho/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8e3843d8b123140dc43dc03cfd06a1cbd6954f9a5a8ec01276d42590b2078b3d
+oid sha256:f5cfceb0f75c031f5021a0fdf52f1fdb2e14b441dbf38a28fd6cb6bf5af2d9bf
 size 748434

butterfly_ortho/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8f8c940840c1f30b31c928ad24cd6f3017e1bb48ae200a6b9e17ccc3928c8eb
+size 476578

butterfly_ortho/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3eaa20b37a33ff9295be89b0d8ba1582f0730919fc75549151368c26a6bc0823
+size 865393

monarch_8/config.json

@@ -24,6 +24,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

monarch_8/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:16eba9c4fec8390ecccc7519a4b87882c34d6aeb2afaa080e3cefcd6b27509a8
+oid sha256:c7416829373b77c1a0fd9ef7a33681131abe33cda0bdf95e8ddc585610a7995b
 size 1423570

monarch_8/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:368f9a6af401f42000294c45d5b827dd275d9fb2c486f2e687943d3f23451251
+size 1465674

monarch_8/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0776227333ef5075c967846daae9e3ecbdc8f924dc0905002a6b5f00703eec49
+size 1442161

monarch_fc/config.json

@@ -23,6 +23,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

monarch_fc/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b80535acab1d3dc1202720a62d10cdb896a5a25aeb65308bd3d05bf4f1448a88
+oid sha256:cde0e25bebea3d5d713afaaf10f5e52eabaeb5ef552de45abc1cb216957e1848
 size 8570351

monarch_fc/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:febb340361f7f9c015f45ebdb031dee822a6e652eed48f18382c140cac878122
+size 2891983

monarch_fc/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bda161d596109d21b794656776329f4986cdd78be7362edafbb2b57534cae16
+size 8578184

monarch_full/config.json

@@ -24,6 +24,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

monarch_full/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74fed92103025cdd7202efd661f3a3e5fa52cf4f018a6df5a9f382f3c372a57f
+oid sha256:c9a34fe8ba46ad1c41ac474010eca2a8bbaadf30d334901af81146b93cc95ba7
 size 2811538

monarch_full/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:640d277c5843df8c4e1b0518e20c29a7f377847a436b62f02ae0ff43e4b6cd04
+size 2853656

monarch_full/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:accb5c6dfc79dde3f7772ed90b67ce9b9dc3fcb06b26de9784d4f0d100099b64
+size 2830166

wide_monarch/config.json

@@ -24,6 +24,10 @@
     "hop_size": 256,
     "win_size": 512,
     "num_workers": 8,
+    "quant": {
+        "enabled": false,
+        "n_calib_utts": 200
+    },
     "dist_config": {
         "dist_backend": "nccl",
         "dist_url": "tcp://localhost:54321",

wide_monarch/g_best

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ad47e3218d21c26b21343dfec9485ad894de7e45a0bbb3d2300492fb648f2462
+oid sha256:e9ad0dda8f30258d8a9dfa42fd900c4e418bc97bdc2a2ed9a6aae709ab6542d7
 size 9432786

wide_monarch/g_best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc5ac41831080f05d188ca695b9b57a77bf2ee3092008be5fd9ce64f5ecfeb58
+size 9458215

wide_monarch/g_best_fp32.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1af01d6864dbb59f459cd25dc77b68befb28d7963f1164dffcb42b300167fa3b
+size 9451516