exacta

Fast, reference-validated losses and metrics for deep learning, on CPU and GPU.

Every project ends up re-implementing Dice, HD95 and friends, and the implementations quietly disagree: Podobnik & Vrtovec (2024) compared 11 open-source tools and found that distance-based metrics such as HD95 and NSD can differ substantially across tools for the same segmentation pair. exacta aims to be the one implementation you can trust and cite: every metric states its exact definition and is tested against independent reference implementations (SciPy, MONAI, MedPy, DeepMind's surface-distance, MeshMetrics, SimpleITK, Panoptica and the BraTS 2023 evaluation code), and where they disagree with each other, exacta can reproduce each of them through an explicit argument.

Status: early development (v0.0.1a2). Current releases target 3D medical image segmentation.

Installation

Install the alpha from PyPI with Python 3.10 or newer:

python -m pip install exacta==0.0.1a2

This release is a source distribution. Installation builds the native extension locally and requires a C++17 compiler and Python development headers; OpenMP is optional. On Debian/Ubuntu, install build-essential and the development headers matching your Python version first. Portable binary wheels are planned.

Usage

import exacta.functional as EF
from exacta.losses import DiceCELoss
from exacta.metrics import HD95, NSD, Dice

# Training
loss_fn = DiceCELoss(include_background=False)        # logits (B, C, D, H, W), labels (B, D, H, W);
                                                      # fused Triton kernels on GPU
loss = loss_fn(model(x), y)

# Validation: accumulate per-case values across batches (and processes under torch.distributed)
dice = Dice(num_classes=4, include_background=False)
hd95 = HD95(num_classes=4, include_background=False)
nsd = NSD(tolerance=[1.0, 2.0, 2.0], num_classes=4, include_background=False)  # mm, per class
for x, y, spacing in val_loader:
    pred = model(x).argmax(1)
    dice.update(pred, y)
    hd95.update(pred, y, spacing=spacing)             # per-case voxel size in mm, (B, 3)
    nsd.update(pred, y, spacing=spacing)

result = hd95.compute()
result.per_case    # (N, C)
result.per_class   # (C,)
result.mean        # scalar

# Stateless functions accept PyTorch, NumPy, CuPy or JAX arrays and return the same type
EF.hd95(pred_np, target_np, spacing=(1.5, 0.8, 0.8), num_classes=4)  # -> np.ndarray (B, C)

# Volume and detection rates: sensitivity, precision, specificity, volume similarity, RVD
EF.sensitivity(pred, target, num_classes=4, zero_division=0.0)          # 0 where the target is empty
EF.volume_similarity(pred, target, num_classes=4, convention="itk")      # SimpleITK's signed variant

# Match another tool's definition explicitly
s = (1.5, 0.8, 0.8)                                                  # voxel size in mm
EF.hd95(pred, target, spacing=s, num_classes=4, combine="pooled")    # MedPy's hd95
EF.assd(pred, target, spacing=s, num_classes=4, combine="mean")      # MedPy 0.4's assd
EF.nsd(pred, target, spacing=s, num_classes=4, tolerance=1.0, surface="surfels")  # DeepMind's NSD
EF.hd95(pred, target, spacing=s, num_classes=4, surface="mesh")                  # MeshMetrics' HD95

# Lesion-wise (instance) evaluation: choose how lesions are matched
from exacta import lesions
scores = lesions.evaluate(pred, target, num_classes=4, include_background=False, connectivity=3,
                          matching="greedy", score="iou", threshold=0.5)   # Panoptica's default
scores.f1, scores.sq_iou, scores.pq_iou, scores.lesion_dice                # each (B, C)
brats = lesions.brats2023(lesions.brats_regions(pred), lesions.brats_regions(target),
                          challenge="GLI", spacing=(1.0, 1.0, 1.0))      # BraTS 2023, WT/TC/ET

# Connected components of one binary image (no batch or channel axis)
from exacta import components
labels, count = components.label(mask, connectivity=3)        # numbered exactly as scipy.ndimage.label
kept = components.largest(mask, connectivity=3)                # e.g. nnU-Net-style post-processing
cleaned = components.remove_small(mask, 50, connectivity=1)    # drop components under 50 voxels

Conventions

Each convention is an explicit argument, because this is where libraries disagree.

Spacing components must lie in [1e-100, 1e100] for voxel surfaces and [1e-50, 1e50] for surfels and meshes. The latter need a narrower range because area calculations and triangle predicates involve fourth powers of physical coordinates. Unsupported values raise ValueError, including for empty masks.

Under torch.distributed, every rank must call compute(), including ranks without any local updates. All ranks receive the gathered scores; if no rank has cases, all raise RuntimeError. Use an evaluation sampler without padding so repeated cases do not bias the average. With NCCL, set each rank's current CUDA device before computing.

Training losses

Losses are PyTorch modules taking floating-point logits (B, C, *spatial) and integer class indices (B, *spatial). Use target_format="one-hot" for masks or soft labels with the logits' shape. Softmax targets should sum to one over classes; sigmoid targets should lie in [0, 1]. Target probabilities are used as supplied, without normalization or clamping. For binary segmentation with a single output channel, use sigmoid and a same-shaped target.

from exacta.losses import DiceLoss, DiceCELoss, TverskyLoss, FocalLoss

dice_ce = DiceCELoss(include_background=False)
tversky = TverskyLoss(alpha=0.3, beta=0.7, include_background=False)
focal = FocalLoss(gamma=2.0)  # softmax, no class weighting
binary_focal = FocalLoss(activation="sigmoid", target_format="one-hot", alpha=0.25)
loss = tversky(logits, labels)
loss.backward()

For probabilities p and targets g, the definitions and reductions are explicit:

Dice and Tversky sum over spatial axes, or batch and spatial axes with batch_dice=True. include_background=False omits channel 0 after activation. Their smooth defaults to 1e-5. Tversky's alpha weights false positives and beta false negatives; both default to 0.5 and need not sum to one. At those defaults, Tversky equals Dice with twice the smoothing constant. The asymmetric overlap objective follows Salehi et al. (2017).

GeneralizedDiceLoss follows Sudre et al. (2017) with weight_type="square" (inverse squared target volume), "simple" (inverse volume), or "uniform" (unit weights). It computes a ratio of weighted sums over classes, rather than DiceLoss's mean of classwise ratios. With batch_dice=True, volumes and weights cover the whole local batch. Smoothing is added after weighting and can materially affect the result when inverse-volume weights are small.

Absent classes receive the largest weight among present selected classes. For inverse-volume weighting, empty_target="zero" follows MONAI: cases with no selected target class contribute zero loss, including when predictions contain false positives. empty_target="uniform" instead assigns unit weights to these cases; use positive smoothing for a useful gradient on an entirely empty target. Zero-weight cases return zero even with smooth=0; other zero denominators remain NaN. Soft targets use the product-based formula above. No background channel is added implicitly for single-channel sigmoid inputs.

from exacta.losses import GeneralizedDiceLoss

loss_fn = GeneralizedDiceLoss(include_background=False, empty_target="uniform")
loss = loss_fn(logits, labels)

Focal loss downweights easy examples using gamma (default 2), following Lin et al. (2017). alpha=None gives unit weights. A scalar alpha in [0, 1] assigns sigmoid positive/negative weights alpha/1-alpha; for softmax, channel 0 gets 1-alpha and every other class gets alpha. With gamma=0, alpha=None, focal loss equals CE/BCE, including its reduction. Softmax focal loss is therefore C times MONAI's default mean for the same hard targets. The formulas above define soft targets as weighted sums of hard-label losses; sigmoid soft-target behavior differs from MONAI's modulation of a single soft-label BCE term. Focal loss expects finite logits and includes every channel.

All losses compute half/bfloat16 inputs in float32 and retain float64 precision. backend="auto" uses fused Triton on CUDA when available and PyTorch elsewhere. Use backend="torch" for higher-order derivatives; the fused backward explicitly rejects them. Targets requiring gradients automatically select PyTorch. Eager execution rejects out-of-range indices; compiled Triton execution returns NaN for them without a host synchronization.

For repeated CPU training steps, torch.compile can reduce allocations and fuse native PyTorch operations. Compilation adds startup cost; benchmark the warmed-up loss on your shapes:

import torch

loss_fn = torch.compile(DiceCELoss(backend="torch"), fullgraph=True)

Architecture

• Compiled CPU kernels (csrc/), multithreaded with OpenMP, for every pass over the voxels that dominates a metric: the exact Euclidean distance transform behind HD95, ASSD and NSD; surface extraction (boundary voxels, neighbour codes, the marching-cubes mesh) and point-to-mesh distances; the voxel counts of the overlap metrics, in one pass that also checks the values; connected-component labelling and the label statistics behind largest, remove_small and lesion matching; dilation; and the sort behind area-weighted percentiles. They sit behind a plain C interface and receive arrays through DLPack, so the extension never links against PyTorch and one wheel works across PyTorch versions. They are built without fused multiply-adds, which would round reference expressions differently.
• GPU kernels in Triton, compiled at run time without the CUDA toolkit. backend="auto" (the default everywhere) picks them on GPU; backend="torch" remains as a fallback.
  • The distance transform runs the same algorithm as the CPU kernel and gives the same values, bit for bit. Labelling gives the same labels as the CPU kernel and SciPy.
  • The overlap metrics count voxels in per-block histograms, checking the values as they go.
  • Dice, Dice-plus-CE, Tversky, generalized Dice and focal losses fuse voxel operations in forward and backward, retaining the inputs and small reduction statistics instead of full-volume activations or one-hot intermediates. They work under torch.compile without graph breaks. CPU losses use PyTorch's compiled operators and autograd; torch.compile can fuse them further.
• PyTorch for the glue between kernels, the CPU path of the losses, and the fallbacks.

Benchmarks

AMD Ryzen 9 5900XT (16 cores), NVIDIA RTX 5070, PyTorch 2.14.1, CUDA 13.0. Times are in milliseconds. Metric timings use the median of 5 runs (one run for anything slower than 5 s). Reproduce them with python benchmarks/bench_metrics.py --sizes 128 256 (add --only exacta to skip the slow references). Loss measurements and commands are described below.

Metrics

One synthetic 3-D prediction/target pair with 4 classes, spacing 1.5 × 0.8 × 0.8 mm. The references run on CPU (MONAI's GPU path needs cuCIM, which is not installed here).

"SciPy EDT" is exacta's test reference (tests/reference.py): the textbook implementation on scipy.ndimage.distance_transform_edt. exacta's plain-PyTorch GPU fallback (backend="torch") takes 3,090 ms for HD95 at 256³. Mesh distances run on the CPU, also for CUDA inputs. The fastest CPU kernels spend much of their time writing fresh output memory, so their times depend on how quickly the operating system supplies it: labelling at 256³ takes 6-13 ms from run to run. The BraTS row scores the regions WT, TC and ET of one synthetic case from the lesion test fixtures' generator (tests/fixtures/make_lesion_cases.py); BraTS-2023-Metrics needs an environment of its own (see below) and was timed there, on the same case.

Losses

The current benchmark uses identical seeded inputs for eager PyTorch, torch.compile over backend="torch", and fused Triton. Compilation and 5 warm-up steps are excluded; results are medians of 20 forward-plus-backward steps with synchronization. Default loss arguments, softmax activation, integer targets, float32 logits of shape 2×4×128³, RTX 5070:

Extra memory is the peak allocation during the step beyond its inputs and logits' gradient; small scratch buffers can fit below that baseline. These are local synthetic measurements, not end-to-end training speedups. CPU measurements on the same system, with 16 threads and 2×4×64³ float32 logits, use the same seed, warm-up and repeat count:

Reproduce these tables with:

python benchmarks/bench_losses.py --classes 4 --sizes 128 --dtypes float32 --repeat 20 --warmup 5 --seed 0
python benchmarks/bench_losses.py --device cpu --classes 4 --sizes 64 --dtypes float32 \
    --backends torch compiled --threads 16 --repeat 20 --warmup 5 --seed 0

The earlier Dice-plus-CE scaling measurements below use the median of 10 warmed-up GPU steps, integer targets and the same peak-memory definition. The benchmark also supports bfloat16, float64, sigmoid activation and other batch, class and volume sizes.

Development

python -m venv ~/.venvs/exacta && source ~/.venvs/exacta/bin/activate
pip install torch scikit-build-core nanobind cmake ninja
pip install --no-build-isolation -e ".[test]"   # re-run after editing csrc/
pytest
python benchmarks/bench_metrics.py
python benchmarks/bench_losses.py

Reference tests against MONAI, MedPy, DeepMind's surface-distance, MeshMetrics, VTK and SimpleITK are skipped unless those packages are installed:

pip install monai medpy scikit-image vtk SimpleITK \
    "git+https://github.com/google-deepmind/surface-distance.git" \
    "git+https://github.com/gasperpodobnik/MeshMetrics.git"

Lesion-wise evaluation is tested against fixtures produced by Panoptica 2.1.7 and by the BraTS 2023 evaluation code, which need environments of their own (old NumPy and pandas). tests/fixtures/regenerate.sh rebuilds those environments with uv, reruns both tools and reports whether their results changed (set PANOPTICA_VERSION or BRATS_COMMIT to try other releases).

Array-library round-trip tests run for NumPy, and for CuPy and JAX when installed: python -m pytest tests/test_inputs.py -k round_trip. Distributed tests exercise ranks without updates using Gloo; their NCCL variants require two GPUs and are skipped otherwise.

Loss checks include independent NumPy formulas, MONAI comparisons with explicit reduction adjustments, float64 finite-difference gradient checks, higher-order PyTorch gradients, CPU/GPU agreement, mixed precision, extreme logits and full-graph compilation: python -m pytest tests/test_losses.py tests/test_weighted_losses.py tests/test_generalized_dice.py.

License

Apache License 2.0; see LICENSE and NOTICE.