[davidackerman]

Stacked on #142 (browser-full-augment), itself on #137. Once those merge, this PR's base becomes main.

The browser can now train the whole network of the 3D models, setup_53, setup_54, setup_55, setup_56, setup_61 and setup_67, not only their last layer. This builds on the export Marwan Zouinkhi added to cellmap/onnx-web-models on 2026-10-17 (3331842ec2): a 3D training graph per folder, torch_params.json and weight_map.json to read the weights out of the served .onnx and write them back, and a reference runtime.

Who sees a change

• The Finetune tab, with a 3D model served: Training now offers Whole network beside Last layer.
  • Tile Size becomes Block Size: 124³ (the default), 140³, 156³ or 172³.
  • The hint gives each size's measured time and memory.
  • Batch Size stays at 1: one block a step.
• The training data: it trains on cubes around the paint, so strokes on neighbouring slices teach it along z too. Affinity targets get their z neighbour.
• The same settings as the 2D whole network: loss type, smoothing, class balance, distillation (the starting network run on the same block), and augmentation. In 3D, augmentation flips z, y and x, and turns in y–x where those voxels are square.
• A block too large for the GPU stops the run with a message naming a smaller size to try. The weights are left as they were.
• 2D models and the last layer are unchanged.

Memory and speed

Marwan's runtime trains with TensorFlow.js's own gradients. Those keep every intermediate: about 8 GB at 124³, and nothing larger fits an 11 GB GPU. As written, its backward pass also exceeds WebGPU's dispatch limit from 108³. net3d.ts builds the same network from the same graph and weights, but trains it lean:

• every op's gradient is written out;
• conv+BatchNorm+ReLU and InstanceNorm(+add)(+ReLU) are fused;
• a convolution's gradients are matrix products;
• the backward pass recomputes segments from checkpoints;
• the GPU is synced and TF.js's buffer pool emptied between segments.

One step with the starting network's forward pass, on an RTX 2080 Ti (TF.js 4.22.0, WebGPU, Chromium):

Model	124³	140³	172³
setup_53 (UNet; setup_55 the same)	2.8 s, 2.2 GB	4.1 s, 3.0 GB	8.2 s, 5.3 GB
setup_54 (ResUNet; 56, 61, 67 the same)	4.0 s, 4.0 GB	5.9 s, 5.5 GB	11.9 s, 9.9 GB

Out of memory raises nothing in TensorFlow.js; it just corrupts the step. Each step therefore runs inside WebGPU error scopes, and Adam's update is applied only when they report nothing. Tested: setup_54 at 196³ threw cleanly, none of its weights changed, and 124³ trained normally afterwards.

Code

• full/cellmap-tfjs-train.ts: Marwan's runtime at 3331842ec2, ported to TypeScript; the 2D path is unchanged. Adds the 3D buildNet, readTorchParams, toOnnxWeights, and a checked interpreter for the structured transforms.
• full/net3d.ts (new): the lean 3D network behind buildNet3d, with step, forward/forwardAsync and exportWeights.
• full/graph.ts: NDHWC graphs and N³ block sides (92 + 8k, context 88).
• full/tiles.ts, full/run-full.ts: 3D blocks around the paint, the raw around each, 3D targets, the 3D run, and write-back through weight_map.json.
• modes.ts: offers the whole network for a 3D model whose manifest names torch_params and a weight map. This lifts Batch: PRs #126–#141 in one merge #137's 2D-only guard (74bbd72).
• Wiring: the routes, form, job card, docs (browser.rst) and CHANGELOG.

Checks

• npm run check in web/: typecheck, 1,558 tests passed, build.
• The full Python suite: 2,230 passed, with only test_importing_leaves_the_process_alone[blockwise-logging] failing, as on main in that environment.
• Tests:
  • Lean vs Marwan's runtime: on tiny UNet and ResUNet graphs, two Adam steps each, the lean network gives the same losses and weights as his plain one. A deliberately wrong gradient makes the test fail.
  • Weights: readTorchParams/toOnnxWeights reproduce the served bytes exactly. An opt-in live test round-trips setup_53's real fp32 and fp16 files byte for byte.
  • Blocks: a 3D block's raw and paint stay registered through flips and turns, and its z-affinities are right.
  • Offering: the routes offer the 3D sides with costs and refuse a batch above 1.
• Agreement: setup_53's logits at 124³ match onnxruntime-web to 2.3e-5, with no sign flips in 139,968. setup_54's match PyTorch to 2.9e-5. setup_54's served ONNX is itself about 6e-3 from PyTorch in onnxruntime.
• End to end, headless on this workstation's GPU, jrc_hela-2 with setup_53 at 124³, through the real form:
  • 3 epochs, with augmentation and distillation;
  • the finetuned model was served and its chunks came back;
  • the download is 22.5 MB.

Not included

• Good regions aren't used by the whole network, 2D or 3D.
• The first steps are slower (7–9 s at 124³) while the viewer is still computing the base model's first chunks. The log's "A step takes" is measured at step 2, so it overstates a 3D step on a fresh page.
• Distillation at λ 0.01 holds the network only lightly: on a synthetic block the term reached 213 by epoch 3. Worth testing on real paint before choosing a default.