| --- |
| library_name: onnx |
| tags: [onnx, electron-microscopy, segmentation, cellmap, webgpu, onnxruntime-web] |
| --- |
| |
| # CellMap ONNX web models |
|
|
| Small 3D U-Nets for FIB-SEM organelle segmentation (mitochondria, nuclei, ER), exported to ONNX for **in-browser** inference with |
| [onnxruntime-web](https://onnxruntime.ai/) in the [cellmap-flow serverless client](https://github.com/mzouink/cellmap-flow/tree/serverless). |
|
|
| Every file takes `input` float32 `[1,1,z,y,x]` and returns `output` float32 **raw logits** `[1,C,z',y',x']` |
| (apply sigmoid, or threshold at 0). Input normalization: `uint8 / 255 * 2 - 1` |
| (`[{"name":"MinMaxNormalizer","min_value":0,"max_value":255},{"name":"LambdaNormalizer","expression":"2*x-1"}]`). |
| The input block size is fixed per file (`_in<N>` = N³ voxels); valid convolutions, so the output is N−88 per axis. |
|
|
| Two graph layouts of the **same models** (same weights, same outputs up to float rounding): |
|
|
| - `setup_XX/…` — **recommended**: every Conv3d rewritten as Conv2d over a depth-folded batch (~16× faster on WebGPU, see below). |
| - `conv3d/setup_XX/…` — native Conv3d nodes, for runtimes with a fast 3D conv (e.g. onnxruntime CUDA/CPU, TensorRT). |
|
|
| ## Models |
|
|
| | folder | target | channels | voxel (nm) | architecture | params | checkpoint | val F1 | |
| |---|---|---|---|---|---|---|---| |
| | `setup_54` | mito_affinities | mito_aff_z, mito_aff_y, mito_aff_x | 16 | ResUNet3DValid | 10.7M | 260000 | 0.958 | |
| | `setup_53` | mito_affinities | mito_aff_z, mito_aff_y, mito_aff_x | 16 | Unet3dIso | 5.6M | 340000 | 0.955 | |
| | `setup_56` | nuc | nuc | 128 | ResUNet3DValid | 10.7M | 350000 | 0.765 | |
| | `setup_55` | nuc | nuc | 128 | Unet3dIso | 5.6M | 330000 | 0.645 | |
| | `setup_61` | nuc | nuc | 256 | ResUNet3DValid | 10.7M | 20000 | 0.738 | |
| | `setup_67` | er | er | 8 | ResUNet3DValid | 10.7M | 20000 | 0.806 | |
|
|
| ## Download links and speed (recommended 2D-rewrite layout) |
|
|
| Times are steady-state ms per block on WebGPU (Chrome, Apple M-series GPU, real EM; see benchmark below). |
| Lower ms is better; `out Mvox/s` = output voxels per second (higher is better). fp16 is 1.3–2.2× faster than fp32, |
| halves download and GPU memory, and changes ≤0.05% of `logit > 0` decisions. fp32 is exact and also runs on the WASM (CPU) fallback. |
| The largest block of each model gives the best throughput (88 voxels of context are lost per axis per block). |
|
|
| | model | target | voxel | in → out | fp32 | fp32 ms | fp16 | fp16 ms | fp16 speedup | out Mvox/s fp32 / fp16 | |
| |---|---|---|---|---|---|---|---|---|---| |
| | `setup_54` | mito_affinities (3 ch) | 16 nm | 148³ → 60³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in148_fp32.onnx) | 1203 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in148_fp16.onnx) | 905 | 1.33× | 0.18 / 0.24 | |
| | `setup_54` | mito_affinities (3 ch) | 16 nm | 132³ → 44³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in132_fp32.onnx) | 796 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in132_fp16.onnx) | 573 | 1.39× | 0.11 / 0.15 | |
| | `setup_54` | mito_affinities (3 ch) | 16 nm | 116³ → 28³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in116_fp32.onnx) | 516 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_54/setup_54_mito_affinities_16nm_in116_fp16.onnx) | 323 | 1.60× | 0.04 / 0.07 | |
| | `setup_53` | mito_affinities (3 ch) | 16 nm | 172³ → 84³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in172_fp32.onnx) | 1137 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in172_fp16.onnx) | 741 | 1.53× | 0.52 / 0.80 | |
| | `setup_53` | mito_affinities (3 ch) | 16 nm | 140³ → 52³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in140_fp32.onnx) | 533 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in140_fp16.onnx) | 336 | 1.59× | 0.26 / 0.42 | |
| | `setup_53` | mito_affinities (3 ch) | 16 nm | 124³ → 36³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in124_fp32.onnx) | 348 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_53/setup_53_mito_affinities_16nm_in124_fp16.onnx) | 157 | 2.22× | 0.13 / 0.30 | |
| | `setup_56` | nuc (1 ch) | 128 nm | 148³ → 60³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in148_fp32.onnx) | 1208 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in148_fp16.onnx) | 921 | 1.31× | 0.18 / 0.23 | |
| | `setup_56` | nuc (1 ch) | 128 nm | 132³ → 44³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in132_fp32.onnx) | 811 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in132_fp16.onnx) | 574 | 1.41× | 0.11 / 0.15 | |
| | `setup_56` | nuc (1 ch) | 128 nm | 116³ → 28³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in116_fp32.onnx) | 527 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_56/setup_56_nuc_128nm_in116_fp16.onnx) | 324 | 1.63× | 0.04 / 0.07 | |
| | `setup_55` | nuc (1 ch) | 128 nm | 172³ → 84³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in172_fp32.onnx) | 1080 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in172_fp16.onnx) | 758 | 1.42× | 0.55 / 0.78 | |
| | `setup_55` | nuc (1 ch) | 128 nm | 140³ → 52³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in140_fp32.onnx) | 506 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in140_fp16.onnx) | 319 | 1.59× | 0.28 / 0.44 | |
| | `setup_55` | nuc (1 ch) | 128 nm | 124³ → 36³ | [fp32 (23 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in124_fp32.onnx) | 331 | [fp16 (11 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_55/setup_55_nuc_128nm_in124_fp16.onnx) | 156 | 2.12× | 0.14 / 0.30 | |
| | `setup_61` | nuc (1 ch) | 256 nm | 148³ → 60³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in148_fp32.onnx) | 1220 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in148_fp16.onnx) | 900 | 1.36× | 0.18 / 0.24 | |
| | `setup_61` | nuc (1 ch) | 256 nm | 132³ → 44³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in132_fp32.onnx) | 809 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in132_fp16.onnx) | 574 | 1.41× | 0.11 / 0.15 | |
| | `setup_61` | nuc (1 ch) | 256 nm | 116³ → 28³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in116_fp32.onnx) | 514 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_61/setup_61_nuc_256nm_in116_fp16.onnx) | 322 | 1.60× | 0.04 / 0.07 | |
| | `setup_67` | er (1 ch) | 8 nm | 148³ → 60³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in148_fp32.onnx) | 1233 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in148_fp16.onnx) | 890 | 1.39× | 0.18 / 0.24 | |
| | `setup_67` | er (1 ch) | 8 nm | 132³ → 44³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in132_fp32.onnx) | 809 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in132_fp16.onnx) | 569 | 1.42× | 0.11 / 0.15 | |
| | `setup_67` | er (1 ch) | 8 nm | 116³ → 28³ | [fp32 (43 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in116_fp32.onnx) | 532 | [fp16 (22 MB)](https://huggingface.co/cellmap/onnx-web-models/resolve/main/setup_67/setup_67_er_8nm_in116_fp16.onnx) | 323 | 1.65× | 0.04 / 0.07 | |
|
|
| Native-3D versions of every file are under `conv3d/` with the same names (e.g. `https://huggingface.co/cellmap/onnx-web-models/resolve/main/conv3d/setup_54/setup_54_mito_affinities_16nm_in148_fp32.onnx`). On WebGPU they are ~16× slower (setup_54 148³ fp32: 20.9 s vs 1.2 s per block) and fp16 does not help there (17–25 s), because onnxruntime-web runs Conv3d with a naive shader; use them for onnxruntime CUDA/CPU or TensorRT. |
| |
| ## Compression variants |
| |
| | variant | weights / compute | file vs fp32 | notes | |
| |---|---|---|---| |
| | `fp32` | fp32 / fp32 | 1× | reference; best accuracy, works on WebGPU and WASM (CPU) | |
| | `fp16` | fp16 / fp16 | ½ | halves download and activation memory; WebGPU only (poor fp16 support on WASM) | |
| | `int8w` | int8 per-channel / fp32 | ¼ | smaller download only (weights dequantized at load) | |
| | `int8w_fp16` | int8 per-channel / fp16 | ¼ | smallest download + fp16 activations | |
|
|
| ## Graph rewrites for WebGPU |
|
|
| Both are exact (same math, same weights; verified against the PyTorch model): |
|
|
| - **Conv3d → Conv2d over a depth-folded batch.** Each 3D conv is computed as a sum of 2D convs, one per depth tap, on the |
| volume reshaped to `[depth, C, H, W]`. onnxruntime-web runs 3D convs with a naive WebGPU shader but 2D convs with tuned |
| matmul kernels, so this is ~16× faster (setup_54, 148³ block, Apple M-series GPU: 20.9 s → 1.2 s). |
| - **3D ConvTranspose → 1×1×1 conv + depth-to-space** (Unet3dIso decoders); the WebGPU EP has no 3D ConvTranspose. |
| |
| Peak memory is dominated by activations, not weights: the largest single tensor is ~40×N³ (ResUNet) or ~32×(N−2)³ (Unet3dIso) |
| values, e.g. 259 MB fp16 / 519 MB fp32 at N=148. Pick a smaller block if your GPU's WebGPU `maxBufferSize` is lower. |
| |
| ## Browser WebGPU benchmark |
| |
| Measured in Chrome 154 on macOS, Apple GPU (WebGPU `metal-3`), onnxruntime-web 1.27.0-dev.20260506; best of 3 runs after one warm-up run. |
| |
| Real EM (jrc_mus-liver, public OpenOrganelle data) at each model's voxel size. `vs fp32 flip %` = voxels whose `logit > 0` |
| decision differs from the same model's fp32 file on WebGPU. |
|
|
| | model | block | variant | ms / block | output voxels / s | max abs vs PyTorch (sampled) | flip % vs fp32 | |
| |---|---|---|---|---|---|---| |
| | `setup_54` | 148³ | `fp32` | 1203 | 179,551 | 0.001 | — | |
| | `setup_54` | 148³ | `fp16` | 905 | 238,674 | 0.082 | 0.0 | |
| | `setup_54` | 148³ | `int8w` | 1216 | 177,632 | 0.106 | 0.0 | |
| | `setup_54` | 148³ | `int8w_fp16` | 879 | 245,734 | 0.153 | 0.0 | |
| | `setup_54` | 132³ | `fp32` | 796 | 107,015 | 0.0 | — | |
| | `setup_54` | 132³ | `fp16` | 573 | 148,663 | 0.072 | 0.0 | |
| | `setup_54` | 132³ | `int8w` | 1096 | 77,723 | 0.096 | 0.0 | |
| | `setup_54` | 132³ | `int8w_fp16` | 558 | 152,659 | 0.09 | 0.0 | |
| | `setup_54` | 116³ | `fp32` | 516 | 42,543 | 0.001 | — | |
| | `setup_54` | 116³ | `fp16` | 323 | 67,963 | 0.048 | 0.0 | |
| | `setup_54` | 116³ | `int8w` | 531 | 41,341 | 0.103 | 0.0 | |
| | `setup_54` | 116³ | `int8w_fp16` | 328 | 66,927 | 0.101 | 0.0 | |
| | `setup_53` | 172³ | `fp32` | 1137 | 521,288 | 0.001 | — | |
| | `setup_53` | 172³ | `fp16` | 741 | 799,870 | 0.071 | 0.0 | |
| | `setup_53` | 172³ | `int8w` | 1013 | 585,098 | 0.208 | 0.0 | |
| | `setup_53` | 172³ | `int8w_fp16` | 737 | 804,212 | 0.166 | 0.0 | |
| | `setup_53` | 140³ | `fp32` | 533 | 263,805 | 0.0 | — | |
| | `setup_53` | 140³ | `fp16` | 336 | 418,476 | 0.076 | 0.0 | |
| | `setup_53` | 140³ | `int8w` | 494 | 284,632 | 0.268 | 0.0 | |
| | `setup_53` | 140³ | `int8w_fp16` | 328 | 428,683 | 0.225 | 0.0 | |
| | `setup_53` | 124³ | `fp32` | 348 | 134,069 | 0.0 | — | |
| | `setup_53` | 124³ | `fp16` | 157 | 297,172 | 0.073 | 0.0 | |
| | `setup_53` | 124³ | `int8w` | 317 | 147,180 | 0.194 | 0.0 | |
| | `setup_53` | 124³ | `int8w_fp16` | 160 | 291,600 | 0.201 | 0.0 | |
| | `setup_56` | 148³ | `fp32` | 1208 | 178,808 | 0.001 | — | |
| | `setup_56` | 148³ | `fp16` | 921 | 234,528 | 0.2 | 0.0 | |
| | `setup_56` | 148³ | `int8w` | 1219 | 177,194 | 0.379 | 0.0 | |
| | `setup_56` | 148³ | `int8w_fp16` | 880 | 245,455 | 0.34 | 0.0 | |
| | `setup_56` | 132³ | `fp32` | 811 | 105,036 | 0.001 | — | |
| | `setup_56` | 132³ | `fp16` | 574 | 148,404 | 0.16 | 0.0 | |
| | `setup_56` | 132³ | `int8w` | 822 | 103,630 | 0.447 | 0.0 | |
| | `setup_56` | 132³ | `int8w_fp16` | 559 | 152,386 | 0.266 | 0.0 | |
| | `setup_56` | 116³ | `fp32` | 527 | 41,655 | 0.001 | — | |
| | `setup_56` | 116³ | `fp16` | 324 | 67,753 | 0.119 | 0.0 | |
| | `setup_56` | 116³ | `int8w` | 526 | 41,734 | 0.342 | 0.0 | |
| | `setup_56` | 116³ | `int8w_fp16` | 326 | 67,337 | 0.332 | 0.0 | |
| | `setup_55` | 172³ | `fp32` | 1080 | 548,800 | 0.001 | — | |
| | `setup_55` | 172³ | `fp16` | 758 | 781,931 | 0.117 | 0.0 | |
| | `setup_55` | 172³ | `int8w` | 989 | 599,296 | 0.558 | 0.0 | |
| | `setup_55` | 172³ | `int8w_fp16` | 736 | 805,304 | 0.65 | 0.0 | |
| | `setup_55` | 140³ | `fp32` | 506 | 277,881 | 0.001 | — | |
| | `setup_55` | 140³ | `fp16` | 319 | 440,777 | 0.131 | 0.0 | |
| | `setup_55` | 140³ | `int8w` | 481 | 292,324 | 0.531 | 0.0 | |
| | `setup_55` | 140³ | `int8w_fp16` | 336 | 418,476 | 0.512 | 0.0 | |
| | `setup_55` | 124³ | `fp32` | 331 | 140,955 | 0.001 | — | |
| | `setup_55` | 124³ | `fp16` | 156 | 299,077 | 0.103 | 0.0 | |
| | `setup_55` | 124³ | `int8w` | 315 | 148,114 | 0.314 | 0.0 | |
| | `setup_55` | 124³ | `int8w_fp16` | 160 | 291,600 | 0.402 | 0.0 | |
| | `setup_61` | 148³ | `fp32` | 1220 | 177,049 | 0.001 | — | |
| | `setup_61` | 148³ | `fp16` | 900 | 240,000 | 0.124 | 0.0019 | |
| | `setup_61` | 148³ | `int8w` | 1224 | 176,471 | 0.328 | 0.0032 | |
| | `setup_61` | 148³ | `int8w_fp16` | 869 | 248,562 | 0.362 | 0.0037 | |
| | `setup_61` | 132³ | `fp32` | 809 | 105,295 | 0.001 | — | |
| | `setup_61` | 132³ | `fp16` | 574 | 148,404 | 0.164 | 0.0 | |
| | `setup_61` | 132³ | `int8w` | 808 | 105,426 | 0.407 | 0.0 | |
| | `setup_61` | 132³ | `int8w_fp16` | 555 | 153,485 | 0.482 | 0.0 | |
| | `setup_61` | 116³ | `fp32` | 514 | 42,708 | 0.001 | — | |
| | `setup_61` | 116³ | `fp16` | 322 | 68,174 | 0.145 | 0.0 | |
| | `setup_61` | 116³ | `int8w` | 525 | 41,813 | 0.303 | 0.0 | |
| | `setup_61` | 116³ | `int8w_fp16` | 327 | 67,131 | 0.434 | 0.0 | |
| | `setup_67` | 148³ | `fp32` | 1233 | 175,182 | 0.0 | — | |
| | `setup_67` | 148³ | `fp16` | 890 | 242,697 | 0.03 | 0.0148 | |
| | `setup_67` | 148³ | `int8w` | 1221 | 176,904 | 0.042 | 0.0282 | |
| | `setup_67` | 148³ | `int8w_fp16` | 870 | 248,276 | 0.048 | 0.0356 | |
| | `setup_67` | 132³ | `fp32` | 809 | 105,295 | 0.0 | — | |
| | `setup_67` | 132³ | `fp16` | 569 | 149,708 | 0.035 | 0.0129 | |
| | `setup_67` | 132³ | `int8w` | 796 | 107,015 | 0.037 | 0.0399 | |
| | `setup_67` | 132³ | `int8w_fp16` | 556 | 153,209 | 0.047 | 0.0434 | |
| | `setup_67` | 116³ | `fp32` | 532 | 41,263 | 0.0 | — | |
| | `setup_67` | 116³ | `fp16` | 323 | 67,963 | 0.039 | 0.0501 | |
| | `setup_67` | 116³ | `int8w` | 529 | 41,497 | 0.043 | 0.0866 | |
| | `setup_67` | 116³ | `int8w_fp16` | 327 | 67,131 | 0.066 | 0.0729 | |
|
|
| ## Fidelity vs. fp32 PyTorch (onnxruntime CPU) |
|
|
| One real EM crop per model (jrc_c-elegans-comma-1 at the model's voxel size). `flip %` = voxels whose `logit > 0` decision |
| differs from PyTorch. A consistency check, not a re-run of F1 validation. |
| |
| | file | MB | in → out | max abs logit diff | mean abs diff | flip % | |
| |---|---|---|---|---|---| |
| | `setup_54/setup_54_mito_affinities_16nm_in148_fp32.onnx` | 42.9 | 148³ → 60³ | 0.0005 | 0.0001 | 0.0000 | |
| | `setup_54/setup_54_mito_affinities_16nm_in148_fp16.onnx` | 21.6 | 148³ → 60³ | 0.0142 | 0.0024 | 0.0009 | |
| | `setup_54/setup_54_mito_affinities_16nm_in148_int8w.onnx` | 10.9 | 148³ → 60³ | 0.1366 | 0.0198 | 0.0094 | |
| | `setup_54/setup_54_mito_affinities_16nm_in148_int8w_fp16.onnx` | 10.9 | 148³ → 60³ | 0.1139 | 0.0191 | 0.0080 | |
| | `setup_54/setup_54_mito_affinities_16nm_in132_fp32.onnx` | 42.9 | 132³ → 44³ | 0.0003 | 0.0000 | 0.0000 | |
| | `setup_54/setup_54_mito_affinities_16nm_in132_fp16.onnx` | 21.6 | 132³ → 44³ | 0.0149 | 0.0025 | 0.0000 | |
| | `setup_54/setup_54_mito_affinities_16nm_in132_int8w.onnx` | 10.9 | 132³ → 44³ | 0.1452 | 0.0203 | 0.0192 | |
| | `setup_54/setup_54_mito_affinities_16nm_in132_int8w_fp16.onnx` | 10.9 | 132³ → 44³ | 0.1202 | 0.0191 | 0.0160 | |
| | `setup_54/setup_54_mito_affinities_16nm_in116_fp32.onnx` | 42.9 | 116³ → 28³ | 0.0001 | 0.0000 | 0.0000 | |
| | `setup_54/setup_54_mito_affinities_16nm_in116_fp16.onnx` | 21.6 | 116³ → 28³ | 0.0134 | 0.0026 | 0.0015 | |
| | `setup_54/setup_54_mito_affinities_16nm_in116_int8w.onnx` | 10.9 | 116³ → 28³ | 0.1302 | 0.0213 | 0.0137 | |
| | `setup_54/setup_54_mito_affinities_16nm_in116_int8w_fp16.onnx` | 10.9 | 116³ → 28³ | 0.1188 | 0.0209 | 0.0076 | |
| | `setup_53/setup_53_mito_affinities_16nm_in172_fp32.onnx` | 22.5 | 172³ → 84³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_53/setup_53_mito_affinities_16nm_in172_fp16.onnx` | 11.3 | 172³ → 84³ | 0.0618 | 0.0072 | 0.0044 | |
| | `setup_53/setup_53_mito_affinities_16nm_in172_int8w.onnx` | 5.8 | 172³ → 84³ | 0.4799 | 0.0488 | 0.0380 | |
| | `setup_53/setup_53_mito_affinities_16nm_in172_int8w_fp16.onnx` | 5.8 | 172³ → 84³ | 0.5517 | 0.0509 | 0.0351 | |
| | `setup_53/setup_53_mito_affinities_16nm_in140_fp32.onnx` | 22.5 | 140³ → 52³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_53/setup_53_mito_affinities_16nm_in140_fp16.onnx` | 11.3 | 140³ → 52³ | 0.0618 | 0.0073 | 0.0009 | |
| | `setup_53/setup_53_mito_affinities_16nm_in140_int8w.onnx` | 5.8 | 140³ → 52³ | 0.3368 | 0.0430 | 0.0085 | |
| | `setup_53/setup_53_mito_affinities_16nm_in140_int8w_fp16.onnx` | 5.8 | 140³ → 52³ | 0.3255 | 0.0451 | 0.0092 | |
| | `setup_53/setup_53_mito_affinities_16nm_in124_fp32.onnx` | 22.5 | 124³ → 36³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_53/setup_53_mito_affinities_16nm_in124_fp16.onnx` | 11.3 | 124³ → 36³ | 0.0396 | 0.0069 | 0.0007 | |
| | `setup_53/setup_53_mito_affinities_16nm_in124_int8w.onnx` | 5.8 | 124³ → 36³ | 0.2672 | 0.0384 | 0.0129 | |
| | `setup_53/setup_53_mito_affinities_16nm_in124_int8w_fp16.onnx` | 5.8 | 124³ → 36³ | 0.2747 | 0.0395 | 0.0157 | |
| | `setup_56/setup_56_nuc_128nm_in148_fp32.onnx` | 42.9 | 148³ → 60³ | 0.0035 | 0.0003 | 0.0005 | |
| | `setup_56/setup_56_nuc_128nm_in148_fp16.onnx` | 21.6 | 148³ → 60³ | 0.0539 | 0.0065 | 0.0069 | |
| | `setup_56/setup_56_nuc_128nm_in148_int8w.onnx` | 10.9 | 148³ → 60³ | 0.5656 | 0.0669 | 0.0755 | |
| | `setup_56/setup_56_nuc_128nm_in148_int8w_fp16.onnx` | 10.9 | 148³ → 60³ | 0.4653 | 0.0600 | 0.0560 | |
| | `setup_56/setup_56_nuc_128nm_in132_fp32.onnx` | 42.9 | 132³ → 44³ | 0.0019 | 0.0002 | 0.0000 | |
| | `setup_56/setup_56_nuc_128nm_in132_fp16.onnx` | 21.6 | 132³ → 44³ | 0.0543 | 0.0064 | 0.0012 | |
| | `setup_56/setup_56_nuc_128nm_in132_int8w.onnx` | 10.9 | 132³ → 44³ | 0.4437 | 0.0661 | 0.0153 | |
| | `setup_56/setup_56_nuc_128nm_in132_int8w_fp16.onnx` | 10.9 | 132³ → 44³ | 0.4508 | 0.0606 | 0.0117 | |
| | `setup_56/setup_56_nuc_128nm_in116_fp32.onnx` | 42.9 | 116³ → 28³ | 0.0004 | 0.0001 | 0.0000 | |
| | `setup_56/setup_56_nuc_128nm_in116_fp16.onnx` | 21.6 | 116³ → 28³ | 0.0405 | 0.0067 | 0.0046 | |
| | `setup_56/setup_56_nuc_128nm_in116_int8w.onnx` | 10.9 | 116³ → 28³ | 0.4412 | 0.0668 | 0.0820 | |
| | `setup_56/setup_56_nuc_128nm_in116_int8w_fp16.onnx` | 10.9 | 116³ → 28³ | 0.3753 | 0.0611 | 0.0638 | |
| | `setup_55/setup_55_nuc_128nm_in172_fp32.onnx` | 22.5 | 172³ → 84³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_55/setup_55_nuc_128nm_in172_fp16.onnx` | 11.3 | 172³ → 84³ | 0.1344 | 0.0191 | 0.0235 | |
| | `setup_55/setup_55_nuc_128nm_in172_int8w.onnx` | 5.8 | 172³ → 84³ | 0.5892 | 0.0901 | 0.0995 | |
| | `setup_55/setup_55_nuc_128nm_in172_int8w_fp16.onnx` | 5.8 | 172³ → 84³ | 0.5302 | 0.0846 | 0.0957 | |
| | `setup_55/setup_55_nuc_128nm_in140_fp32.onnx` | 22.5 | 140³ → 52³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_55/setup_55_nuc_128nm_in140_fp16.onnx` | 11.3 | 140³ → 52³ | 0.0801 | 0.0183 | 0.0036 | |
| | `setup_55/setup_55_nuc_128nm_in140_int8w.onnx` | 5.8 | 140³ → 52³ | 0.4925 | 0.0885 | 0.0220 | |
| | `setup_55/setup_55_nuc_128nm_in140_int8w_fp16.onnx` | 5.8 | 140³ → 52³ | 0.4770 | 0.0798 | 0.0171 | |
| | `setup_55/setup_55_nuc_128nm_in124_fp32.onnx` | 22.5 | 124³ → 36³ | 0.0000 | 0.0000 | 0.0000 | |
| | `setup_55/setup_55_nuc_128nm_in124_fp16.onnx` | 11.3 | 124³ → 36³ | 0.0660 | 0.0168 | 0.0021 | |
| | `setup_55/setup_55_nuc_128nm_in124_int8w.onnx` | 5.8 | 124³ → 36³ | 0.4925 | 0.0864 | 0.0000 | |
| | `setup_55/setup_55_nuc_128nm_in124_int8w_fp16.onnx` | 5.8 | 124³ → 36³ | 0.4770 | 0.0763 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in148_fp32.onnx` | 42.9 | 148³ → 60³ | 0.0541 | 0.0053 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in148_fp16.onnx` | 21.6 | 148³ → 60³ | 0.0977 | 0.0085 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in148_int8w.onnx` | 10.9 | 148³ → 60³ | 1.2152 | 0.0847 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in148_int8w_fp16.onnx` | 10.9 | 148³ → 60³ | 1.4691 | 0.0912 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in132_fp32.onnx` | 42.9 | 132³ → 44³ | 0.0133 | 0.0015 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in132_fp16.onnx` | 21.6 | 132³ → 44³ | 0.0425 | 0.0061 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in132_int8w.onnx` | 10.9 | 132³ → 44³ | 0.7255 | 0.0742 | 0.0012 | |
| | `setup_61/setup_61_nuc_256nm_in132_int8w_fp16.onnx` | 10.9 | 132³ → 44³ | 0.8130 | 0.0792 | 0.0012 | |
| | `setup_61/setup_61_nuc_256nm_in116_fp32.onnx` | 42.9 | 116³ → 28³ | 0.0023 | 0.0003 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in116_fp16.onnx` | 21.6 | 116³ → 28³ | 0.0283 | 0.0052 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in116_int8w.onnx` | 10.9 | 116³ → 28³ | 0.4741 | 0.0666 | 0.0000 | |
| | `setup_61/setup_61_nuc_256nm_in116_int8w_fp16.onnx` | 10.9 | 116³ → 28³ | 0.4942 | 0.0685 | 0.0000 | |
| | `setup_67/setup_67_er_8nm_in148_fp32.onnx` | 42.9 | 148³ → 60³ | 0.0004 | 0.0000 | 0.0000 | |
| | `setup_67/setup_67_er_8nm_in148_fp16.onnx` | 21.6 | 148³ → 60³ | 0.0127 | 0.0017 | 0.0019 | |
| | `setup_67/setup_67_er_8nm_in148_int8w.onnx` | 10.9 | 148³ → 60³ | 0.1386 | 0.0136 | 0.0171 | |
| | `setup_67/setup_67_er_8nm_in148_int8w_fp16.onnx` | 10.9 | 148³ → 60³ | 0.1204 | 0.0131 | 0.0153 | |
| | `setup_67/setup_67_er_8nm_in132_fp32.onnx` | 42.9 | 132³ → 44³ | 0.0003 | 0.0000 | 0.0000 | |
| | `setup_67/setup_67_er_8nm_in132_fp16.onnx` | 21.6 | 132³ → 44³ | 0.0091 | 0.0020 | 0.0000 | |
| | `setup_67/setup_67_er_8nm_in132_int8w.onnx` | 10.9 | 132³ → 44³ | 0.0754 | 0.0144 | 0.0047 | |
| | `setup_67/setup_67_er_8nm_in132_int8w_fp16.onnx` | 10.9 | 132³ → 44³ | 0.0841 | 0.0145 | 0.0012 | |
| | `setup_67/setup_67_er_8nm_in116_fp32.onnx` | 42.9 | 116³ → 28³ | 0.0001 | 0.0000 | 0.0000 | |
| | `setup_67/setup_67_er_8nm_in116_fp16.onnx` | 21.6 | 116³ → 28³ | 0.0090 | 0.0022 | 0.0364 | |
| | `setup_67/setup_67_er_8nm_in116_int8w.onnx` | 10.9 | 116³ → 28³ | 0.0902 | 0.0194 | 0.2870 | |
| | `setup_67/setup_67_er_8nm_in116_int8w_fp16.onnx` | 10.9 | 116³ → 28³ | 0.0838 | 0.0206 | 0.2824 | |
|
|
| ## Finetuning in the browser and PyTorch originals |
|
|
| Each folder's `manifest.json` has a `finetune` block listing the modes the model offers; a model without the block, or |
| without a mode, is not offered that mode. `source_repo` points to the PyTorch twin (`model_state_dict.pt` + module code |
| `model_def.py`, plus `model.ts` / `metadata.json` for the cellmap-flow server catalog); 2D students also name their |
| 3D teacher in `distilled_from_repo`. |
|
|
| - **head**: train the final 1×1 conv on frozen features (`features` tensor, `weight` / `bias` initializers). |
| - **full**: train the whole network with TF.js on WebGPU: `finetune/train_graph.json` (op list, NHWC for the 2D students, |
| NDHWC for the 3D models) + the reference runtime `finetune/cellmap_tfjs_train.js`. 2D: parameters named as the ONNX |
| initializers (`weight_map` identity). 3D: parameters are the PyTorch state_dict (Conv3d `[O,I,kd,kh,kw]`, |
| ConvTranspose3d `[I,O,2,2,2]`), traced from a native-Conv3d export (BatchNorm unfused) and checked by replaying the graph |
| in PyTorch against the original model. |
| |
| `finetune/initializers.json` gives every initializer's byte offset in each `.onnx` (raw_data inline, fixed shapes), so trained |
| weights are written back in place (as fp16 in the fp16 files) and served as before; int8w files are not writable. |
| In the 3D folders the served graph is rewritten for onnxruntime-web (Conv3d split into `convs.<k>` Conv2d per depth tap, |
| ConvTranspose3d k=s=2 as a 1×1 conv to Cout·8 channels + depth-to-space, channel `c·8 + a·4 + b·2 + e`), so: |
| `finetune/torch_params.json` rebuilds each PyTorch param from the `.onnx` (`from_onnx`), and `finetune/weight_map.json` maps |
| each ONNX initializer back from its PyTorch param (`to_onnx`). Both are JSON op lists — `index [axis,k]`, `permute`, |
| `reshape`, `repeat_interleave [axis,n]`, `stack` — e.g. ConvTranspose3d weight `permute [1,2,3,4,0]` + `reshape [Cout·8,Cin,1,1]`, |
| bias `repeat_interleave [0,8]`; the older Python-syntax `transform` string is kept. |
|
|
| | folder | source_repo | modes | full: step at 556² / 300² (peak) | |
| |---|---|---|---| |
| | `setup_54` | [cellmap/mito-aff-resunet-setup-54](https://huggingface.co/cellmap/mito-aff-resunet-setup-54) | head, full | — | |
| | `setup_53` | [cellmap/mito-aff-unetiso-setup-53](https://huggingface.co/cellmap/mito-aff-unetiso-setup-53) | head, full | — | |
| | `setup_56` | [cellmap/nuc-resunet-setup-56](https://huggingface.co/cellmap/nuc-resunet-setup-56) | head, full | — | |
| | `setup_55` | [cellmap/nuc-unetiso-setup-55](https://huggingface.co/cellmap/nuc-unetiso-setup-55) | head, full | — | |
| | `setup_61` | [cellmap/nuc-resunet-setup-61](https://huggingface.co/cellmap/nuc-resunet-setup-61) | head, full | — | |
| | `setup_67` | [cellmap/er-resunet-setup-67](https://huggingface.co/cellmap/er-resunet-setup-67) | head, full | — | |
| | `setup_53_2d` | [cellmap/mito-aff-unetiso-setup-53-2d](https://huggingface.co/cellmap/mito-aff-unetiso-setup-53-2d) | head, full | 1189 ms (3.9 GB) / 308 ms (1.0 GB) | |
| | `setup_54_2d` | [cellmap/mito-aff-resunet-setup-54-2d](https://huggingface.co/cellmap/mito-aff-resunet-setup-54-2d) | head, full | 1963 ms (10.7 GB) / 490 ms (2.7 GB) | |
|
|
| Full-mode step times (2D students): TF.js 4.22 WebGPU, Chrome, Apple M-series GPU, batch 1, fp32, masked BCE + Adam; peak = TF.js tracked tensors (`manifest.json` → `finetune.full.spike`). No timings published for the 3D models: with plain TF.js autodiff (as in the reference runtime) every intermediate is kept until the backward pass, so only small blocks fit. |
|
|
|
|
| ## Use in cellmap-flow serverless |
|
|
| ONNX model → remote URL `https://huggingface.co/cellmap/onnx-web-models/resolve/main/<folder>/<file>.onnx` |
| (e.g. `setup_54/setup_54_mito_affinities_16nm_in148_fp32.onnx`); |
| Input size `N,N,N`, Output size `N-88,N-88,N-88`, input/output voxel = the model's voxel size, Channels = number of channels |
| (Has channel axis = true for the 3-channel affinity models). Each folder's `manifest.json` has the full metadata. |
| Anonymous Hugging Face downloads are rate-limited (HTTP 429), so clients should cache the model bytes. |
|
|
