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.1a1). The first release targets 3D medical image segmentation.

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

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

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, so installing exacta needs no CUDA toolkit or host compiler. 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 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.

Publishing

The initial alpha publishes the working library as a source distribution. Installation builds the native extension locally and requires a C++17 compiler and Python development headers; OpenMP is optional. Portable binary wheels are still to come. On Debian/Ubuntu, install build-essential and the development headers matching your Python version before installing the package. After the first successful upload, install this alpha with python -m pip install exacta==0.0.1a1.

.github/workflows/publish.yml is triggered manually, only from main. It builds the source archive, checks its metadata and requested version, installs that artifact with PyTorch 2.4.1 and NumPy 1.x, and runs the tests before publishing. Publishing uses short-lived OIDC credentials; no PyPI password or API-token repository secret is needed. Follow PyPI's Trusted Publishing guide.

The account owner must configure the external services once:

1. In GitHub repository settings, create an environment named pypi, restrict deployment to main, and add yourself as a required reviewer. Allow yourself to approve the run.
2. Under PyPI account Publishing, add a pending GitHub publisher with project exacta, owner rodolfo-viana, repository exacta, workflow publish.yml, and environment pypi.
3. In GitHub Actions → Publish to PyPI → Run workflow, select main and enter 0.0.1a1. Review the build/test result and approve the pypi deployment when prompted.

The name is claimed only after the first upload succeeds. Registering a pending publisher alone does not reserve it. For subsequent releases, update __version__ and the workflow's default version, commit the change, then repeat the manual run with that exact version.

License

Apache License 2.0; see LICENSE and NOTICE.