gpu-waterz
Connectomics pipelines consist of:
1. inference on EM volumes (electron-microscopy image stacks of tissue, e.g. brain at nanoscale resolution): a CNN predicts affinities on voxel faces (how likely (==probability of) two neighboring voxels being the same cell);
2. segmentation (decode affinities to labels): watershed to form superpixels, then merge neighbors while mean affinity on the shared boundary stays above a threshold, so every voxel gets an integer segment ID (which cell or fragment it belongs to).
Usually the second step is done by the algorithms (mainly watershed) implemented in the waterz library (CPU based). gpu-waterz is step 2 on GPU: CUDA watershed and contact-mean agglomeration, numpy and torch API; it is not a CPU fallback product.
Watershed, in one sentence intuitively: each voxel “rolls” to whichever neighbor is most strongly predicted to be the same cell, and keeps going until it lands in a small basin; you end up with a huge number of tiny fragments, not finished neurons. More precisely, on the voxel grid you have an affinity on every face between neighbors; watershed builds a directed forest by sending each voxel along the highest-affinity outgoing face (greedy ascent on that field) until you hit a cycle or a local plateau. All voxels that drain to the same attractor share one seed label. Weak faces act as dams. That oversegmentation is deliberate: the later merge step is what fuses those seeds into segments that match the biology.
CREMI is the standard MICCAI benchmark for neuron segmentation from EM: public volumes with expert ground-truth labels. I benchmark on sample A (validation crop), tiled to 2.16 billion voxels at affinity threshold 0.3. Quality is measured with VOI (variation of information), an information-theoretic distance between two segmentations (lower is better). It splits into a split term (extra fragments vs ground truth) and a merge term (distinct cells wrongly fused); connectomics papers usually report both. On 1× RTX 5090 the full gpu-waterz decode is about 3.1 s (affinities already in VRAM). Stock waterz.agglomerate on the same HDF5 is about 9.4 min on the CPU, roughly 180× wall clock on the same machine. Split and merge VOI match stock waterz on that CREMI-A setup (fragment IDs do not need to match).
The slow part of waterz is not the watershed. Step 2 is a priority queue on region-graph edges: take the pair with the highest contact-mean affinity on their shared face, merge them, recompute every neighbor mean on the contracted graph, repeat. That loop is serial; one GPU thread per pop still waits on the last update. Many GPU segmenters dodge that by freezing the original voxel edges and never refreshing means after a merge (Kruskal / mutex / GASP-style pipelines). They can be fast, but they are not the same objective as waterz, and on the same CREMI-A crop they do not reach the split/merge VOI that stock waterz gets.
I kept waterz’s contact-mean statistic and thresholding, but not the heap order. Instead merges use ParHAC (Parallel Hierarchical Agglomerative Clustering, NeurIPS 2022): a parallel substitute for “always merge the current best edge.” Each round builds a matching on the region graph, a set of edges with no shared fragment, so those unions cannot conflict. An edge is eligible if it is (1+ε)-heavy: its contact-mean is within a factor (1+ε) of the strongest edge still incident on each endpoint (the serial heap would only take the global max one edge at a time). All matched pairs merge in parallel; then contact-means are recomputed on the contracted graph, same as after one heap step, and the next round starts. Small ε ≈ fewer, safer parallel merges; large ε ≈ more throughput but you may fuse pairs the heap would not have touched yet, and merge mistakes tend to hurt more than leaving an extra split. The ε used in the repo still matches stock waterz on CREMI-A split/merge VOI (fragment IDs need not match). Float atomicAdd on face sums once made identical affinities give different labels run to run; integer accumulation fixed that so reruns are byte-identical.