BatchAug

Batched GPU augmentations for 3D medical imaging data. The API mirrors MONAI but performs augmentations over an entire batch at once, sampling independent random parameters per batch element (similar to Kornia for 2D). Like MONAI, dictionary transforms apply the same augmentation to volumes and segmentations, keeping paired data aligned. The package provides both PyTorch and Triton backends and automatically selects the fastest one for each transform.

See docs/gallery.md for per-transform visualizations with the init code used to generate each one.

• MONAI-compatible API — drop-in replacements with matching output when B=1 (a few transforms default to recommended fixes; one argument restores exact MONAI parity)
• Independent augmentation across batch (B dimension) — each batch element samples its own random parameters
• Same augmentation across channels (C dimension) — all paired volumes get the same transform
• GPU-native — all operations stay on CUDA, no CPU roundtrips
• Fast — 2–1000x faster than MONAI per-sample loops depending on the transform (benchmarks)
• Auto backend selection — Triton fused kernels where faster, PyTorch/cuDNN elsewhere
• dtype support — works with both float32 and bfloat16

Comparison with Other Libraries

Library	2D	3D	Batched1	GPU
torchvision	✅	❌	❌	○
kornia	✅	○	✅	✅
batchgenerators	✅	✅	✅	❌
TorchIO	✅	✅	❌	✅
monai	✅	✅	❌	○
batchaug	❌	✅	✅	✅

○ = partial support

¹ accepts a batch of examples as input with shape B C H W D and samples augmentation parameters independently for each example

Installation

Requires PyTorch with CUDA. Install PyTorch first following pytorch.org, then install from source:

git clone https://github.com/halleewong/batchaug.git
cd batchaug
pip install -e .

To also install test dependencies:

pip install -e ".[test]"

Usage

All inputs have shape (B, C, H, W, D) where B is the batch size and C is the number of channels (e.g. query/support slices in a few-shot task). The same augmentation parameters are applied to all channels within a batch element, but different parameters are sampled for each batch element.

Task Augmentation

Sometimes it is useful to apply the same augmentations across a set of volumes and segmentations from the same dataset — i.e., task augmentations as in UniverSeg, Tyche, and MultiverSeg. In this setting, the channel dimension stores different examples from the same task.

BatchAug samples parameters independently for each entry in the batch, then applies the same transformation to every entry along the channel dimension. Dictionary transforms apply the same augmentation to all keys, so paired data (e.g. vol + seg) stays aligned. For more details on parameter sampling, see docs/sampling.md.

import torch
import batchaug

# Paired volume and segmentation on GPU
batch = {
    "vol": torch.randn(5, 4, 128, 128, 128, device="cuda"),
    "seg": torch.randn(5, 4, 128, 128, 128, device="cuda"),
}

# Compose a pipeline
task_augs = batchaug.Compose(
    transforms=[
        batchaug.RandRotate90d(keys=["vol", "seg"], prob=0.5, max_k=3, spatial_axes=(0, 1)),
        batchaug.RandAxisFlipd(keys=["vol", "seg"], prob=0.5),
        batchaug.RandGaussianNoised(keys=["vol"], prob=0.5, mean=0.0, std=0.5),
        batchaug.RandAffined(keys=["vol", "seg"], prob=0.5,
                             rotate_range=0.785, shear_range=0.3, translate_range=5),
        batchaug.ScaleIntensityd(keys=["vol"]),
    ],
    lazy=True,
    mode={"vol": "bilinear", "seg": "nearest"},
)

augmented_batch = task_augs(batch)

With lazy=True, geometric transforms (rotations, flips, affines) are fused into a single grid_sample call, avoiding redundant interpolation. Intensity transforms are applied eagerly at their position in the pipeline.

This replaces the slow per-sample loop required by MONAI:

# Before: MONAI per-sample loop (slow, sequential)
monai_aug = monai.transforms.Compose([...])  # same transforms, MONAI API

augmented = {"vol": [], "seg": []}
for i in range(B):
    sample = {"vol": batch["vol"][i], "seg": batch["seg"][i]}
    out = monai_aug(sample)
    augmented["vol"].append(out["vol"])
    augmented["seg"].append(out["seg"])
augmented = {k: torch.stack(v) for k, v in augmented.items()}

The same parameters are used for each channel within a batch element. To perform data augmentation independently across channels, simply reshape the data to merge the B and C dimensions (into the B dimension), apply the augmentation, then reshape back:

# For independent channel augmentation, merge B and C dims
B, C = vol.shape[:2]
vol = vol.view(B * C, 1, *vol.shape[2:])  # (B*C, 1, H, W, D)
aug_vol = t(vol)
aug_vol = aug_vol.view(B, C, *aug_vol.shape[2:])

Data Augmentation

Apply transforms directly to tensors (without dictionary wrapping) when you don't need paired augmentation across keys.

import torch
import batchaug

vol = torch.randn(8, 1, 64, 64, 64, device="cuda")

# Individual transforms
t = batchaug.RandGaussianNoise(prob=0.5, mean=0.0, std=(0.1, 0.5))
noisy_vol = t(vol)

# Or use sample_params / apply for full control
params = t.sample_params(vol.shape[0], vol.shape, vol.device)
noisy_vol = t.apply(vol, params)

Backends: PyTorch and Triton

BatchAug includes two backends: a pure PyTorch backend and a Triton backend with custom fused kernels. When Triton is installed, the library auto-selects the fastest implementation for each transform:

Transform	Default backend	Why
ScaleIntensity	Triton	Fused min/max reduction + rescale (1.5–4.4x faster)
RandAdjustContrast	Triton	Fused normalize + pow + denormalize (1.2–3.0x faster)
RandBiasField	Triton	On-the-fly Legendre polynomial eval avoids large basis tensor (3–10x faster)
RandGaussianSmooth	PyTorch	cuDNN's conv3d is faster than custom Triton separable conv
RandGaussianSharpen	PyTorch	Uses smooth internally, same cuDNN advantage
RandConv	PyTorch	Grouped conv trick is already a single cuDNN call; Triton not expected to help
All others	PyTorch	Already use optimized CUDA ops (cuFFT, grid_sample, etc.)

This happens transparently — batchaug.ScaleIntensity resolves to the Triton version, while batchaug.RandGaussianSmooth resolves to the PyTorch version. You can override:

import batchaug

# Force a specific backend
batchaug.set_backend("pytorch")  # always use PyTorch
batchaug.set_backend("triton")   # always use Triton
batchaug.set_backend("auto")     # auto-select (default)

# Or import from a specific backend directly
from batchaug.pytorch import ScaleIntensity   # PyTorch version
from batchaug.triton import ScaleIntensity    # Triton version

Benchmarks

Per-transform speedup over MONAI's per-sample loop on GPU, using the default auto backend (Triton where faster, PyTorch elsewhere). Measured on NVIDIA L40S with B=5, C=4, 128^3:

Transform	MONAI (ms)	BatchAug (ms)	Speedup
RandBiasField	660.5	0.6	1118x
RandGaussianNoise	744.0	2.6	282x
Rand3DElastic	223.2	9.9	23x
RandAffine	113.5	11.5	10x
RandGibbsNoise	49.0	12.4	4x
ScaleIntensity	2.6	0.9	3x
RandGaussianSmooth	5.8	3.6	2x

The Triton backend provides additional speedups over the PyTorch backend for select transforms:

Transform	Triton vs PyTorch	Why
RandBiasField	3–10x	On-the-fly Legendre eval avoids large basis tensor
ScaleIntensity	1.5–4.4x	Fused min/max reduction + rescale
RandAdjustContrast	1.2–3.0x	Fused normalize + pow + denormalize

Available Transforms

See docs/gallery.md for visualizations of each transform on example data (brain MRI).

Composition

Transform	Description
Compose	Sequential pipeline with optional lazy geometric fusion

Geometric

Transform	Dict version	Description
Rand3DElastic	Rand3DElasticd	Random elastic deformation via smoothed displacement fields
RandAffine	RandAffined	Random affine (rotate, shear, translate, scale) with per-key interpolation modes
RandAxisFlip	RandAxisFlipd	Random flip along a spatial axis
RandFlip	RandFlipd	Flip along specified axes
RandRotate	RandRotated	Arbitrary-angle rotation
RandRotate90	RandRotate90d	Random 90-degree rotation
RandZoom	RandZoomd	Random zoom/scale

Intensity

Transform	Dict version	Description
RandAdjustContrast	RandAdjustContrastd	Gamma correction
RandBiasField	RandBiasFieldd	Polynomial bias field
RandConv	RandConvd	Random convolution (texture/colour perturbation via Xu et al., ICLR 2021)
RandGaussianNoise	RandGaussianNoised	Additive Gaussian noise with per-element mean/std
RandGaussianSharpen	RandGaussianSharpend	Unsharp masking
RandGaussianSmooth	RandGaussianSmoothd	Separable Gaussian blur
RandGibbsNoise	RandGibbsNoised	FFT-based Gibbs ringing
RandRicianNoise	RandRicianNoised	Rician-distributed noise (MRI-specific)
RandScaleIntensity	RandScaleIntensityd	Multiply by random factor
RandScaleIntensityFixedMean	RandScaleIntensityFixedMeand	Scale while preserving mean
RandShiftIntensity	RandShiftIntensityd	Add random offset
RandSimulateLowResolution	RandSimulateLowResolutiond	Downsample/upsample simulation
RandStdShiftIntensity	RandStdShiftIntensityd	Shift by random multiple of per-sample std
ScaleIntensity	ScaleIntensityd	Rescale intensity to [minv, maxv], per element or per element x channel

Utility

Transform	Dict version	Description
DivisiblePad	DivisiblePadd	Pad spatial dims to be divisible by k

Recommended Defaults vs MONAI Parity

BatchAug's API mirrors MONAI, and for almost every transform the defaults match MONAI exactly. For a few transforms, though, we believe MONAI's defaults have a subtle bug or undesirable behavior, so BatchAug defaults to our recommended fix instead of MONAI's value. Each of these is a single argument away from exact MONAI parity — the table below lists the override to pass if you need bit-for-bit MONAI behavior.

Transform	Argument	BatchAug default	MONAI parity
RandSimulateLowResolution	downsample_mode	"nearest-exact"	"nearest"
RandGaussianSmooth	padding_mode	"reflect"	None
RandGaussianSharpen	padding_mode	"reflect"	None

RandSimulateLowResolution — spatial shift from convention mismatch

MONAI's downsample_mode="nearest" paired with upsample_mode="trilinear" shifts the image by ~½ pixel on every down/up roundtrip because "nearest" uses corner-aligned indexing while "trilinear" is cell-centered:

BatchAug defaults to downsample_mode="nearest-exact", which is also cell-centered so the two steps compose cleanly — it keeps the blocky "nearest" look while bringing the shift to <0.02 px (zoom 0.6–0.9):

Pass downsample_mode="nearest" to reproduce MONAI exactly. Other shift-free options are downsample_mode="trilinear" or "area" (true cell-centered resamples, ~0 shift at all zooms, but smoother — no longer a true "nearest" downsample). See docs/defaults.md for a full shift table.

RandGaussianSmooth / RandGaussianSharpen — dark edge halo from zero-padding

MONAI convolves with implicit zero-padding, so voxels within kernel_size // 2 of a boundary average in zeros and pick up a dark halo (and an inverted bright halo on the unsharp-mask sharpen). BatchAug defaults to padding_mode="reflect", which mirrors the volume across each boundary and removes the halo cleanly. The padding_mode argument also accepts "replicate", "constant", or None (implicit zero-padding = MONAI behavior).

Pass padding_mode=None to reproduce MONAI exactly.

Development

Run Tests

python -m pytest tests/ -v