j0126 affinity model (supervised, from scratch)

3-channel voxel affinity prediction for the j0126 zebra finch EM volume, for use with PyTorch Connectomics (tutorials/neuron_j0126).

Files	affinity_scratch_48x96x96.ckpt (247 MB) · affinity_scratch_48x96x96.onnx (247 MB, opset 17)
Architecture	MedNeXt-L, kernel 3 (61.8 M parameters)
Input / output	1 channel EM → 6 channel affinity (banis+; channels 0–2 are the nearest-neighbour affinity used downstream)
Training patch	[48, 96, 96] (ZYX) — anisotropic, ≈ isotropic in nm for the 9×9×20 nm grid
Training data	the 33 dense-GT j0126 cubes, heavy aug_em_neuron augmentation
Schedule	200 k steps from scratch (no pretrained init), AdamW, lr 1e-3, batch 8 × 4 GPUs

This model is not ground-truth-free. It was trained on labelled j0126 tissue and is the supervised upper reference for the zero-shot pipeline in the tutorial, not a drop-in replacement for it. (The dense cubes come from a different alignment than the 50 evaluation skeletons, so it is not voxel-level test contamination — but it is supervision on the target domain.)

Inference window must be [48, 96, 96]

MedNeXt normalizes with GroupNorm(num_groups=C, num_channels=C) — per-sample, per-channel instance norm with no running statistics — so normalization statistics are computed over the sliding window's spatial extent at every block and the forward pass is window-size dependent. Running this checkpoint at the tutorial's inherited [144, 144, 144] is a 3× Z / 1.5× XY mismatch that inverts the trained Z-thin anisotropy.

tutorials/neuron_j0126/1_affinity_supervised.yaml already sets both model.input_size and inference.window.window_size to [48, 96, 96]. Both must be set in YAML, never on the CLI: the stage merge can clobber a pre-resolution CLI override.

Usage

hf download pytc/j0126 affinity_scratch_48x96x96.ckpt --local-dir ckpt/

python scripts/main.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml \
  --mode test --checkpoint ckpt/affinity_scratch_48x96x96.ckpt

Whole volume (726 chunks of 1008³ with a 72-voxel halo), one GPU per shard:

python scripts/main.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml \
  --mode test --checkpoint ckpt/affinity_scratch_48x96x96.ckpt \
  --shard-id "$SLURM_ARRAY_TASK_ID" --num-shards 726

Output is one float16 CZYX HDF5 per chunk. Affinities are written with the scale_sigmoid activation, which the tutorial's ABISS watershed thresholds are calibrated against (scale_sigmoid 0.70 == plain sigmoid 0.9857) — do not substitute a plain sigmoid without recalibrating step 2.

Results

Full-volume j0126, 50 evaluation skeletons, from tutorials/neuron_j0126/README.md. This checkpoint is the "scratch" affinity in that table:

Affinity	Decoding	Error correction	NERL mt=0 ↑	NERL mt=5 ↑	VOI ↓
FFN reference	—	—	0.526	0.538	1.856
scratch	ABISS, exclusion mask	—	0.268	0.470	2.584
scratch	+ nucleus instance certificate	—	0.287	0.482	2.562
scratch	+ nucleus instance certificate	morphology-guided branch linking	0.301	0.539	2.374
scratch	+ nucleus instance certificate	+ 3×3×3 inter-object erosion	0.441	0.528	2.440

Affinity quality alone is not what these numbers measure — see the tutorial for the decode and error-correction steps that produce them.

ONNX

affinity_scratch_48x96x96.onnx is the same weights exported for runtimes without PyTorch.

import numpy as np, onnxruntime as ort

session = ort.InferenceSession("affinity_scratch_48x96x96.onnx",
                               providers=["CPUExecutionProvider"])
image = np.random.randn(1, 1, 48, 96, 96).astype(np.float32)   # (B, C, Z, Y, X)
logits = session.run(["logits"], {"image": image})[0]          # (B, 6, 48, 96, 96)
affinity = 1.0 / (1.0 + np.exp(-0.2 * logits))                 # scale_sigmoid, temperature 0.2

Three things the graph does not do for you:

• The activation. The output is raw logits. The pipeline applies scale_sigmoid with BANIS' temperature 0.2 — sigmoid(0.2 * x) — downstream, and step 2's watershed thresholds are calibrated against that.
• Channel selection. Six channels are emitted; the decode consumes channels 0–2.
• Sliding-window assembly. One window in, one window out. Tiling, overlap and blending are the caller's job.

The spatial dims are fixed at [48, 96, 96] and only the batch axis is dynamic — see the window section above for why a dynamic spatial axis would be a trap here. Deep supervision is switched off in the export, so the graph has a single output.

Verified against PyTorch on identical input: max |diff| 2.2e-04 on logits spanning [1.03, 21.17] (relative 1.0e-05), which is 4.6e-06 after scale_sigmoid — float32 accumulation noise, not a behavioural difference. Regenerate with python scripts/export_onnx.py --config tutorials/neuron_j0126/1_affinity_supervised.yaml --checkpoint <ckpt> --output model.onnx --check.

Checkpoint format

A PyTorch Lightning checkpoint (state_dict keys are prefixed model.model.). Optimizer state has been stripped for download size; the weights are bit-identical to the reference run's step=00200000.ckpt (all 517 tensors, float32). It is loadable for inference/finetuning, not for resuming the original optimizer trajectory.