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).

File	affinity_scratch_48x96x96.ckpt (247 MB)
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.

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.