Batch Augmentations

Use albumentations.batch once your PyTorch DataLoader has assembled the batch. MixUp blends examples and gives you one plan to reuse for their targets. InputNormalizer prepares image or volume tensors on the device where your model runs.

[Dataset: one sample and its targets]
                  |
                  v
[Compose: crop, flip, or adjust the sample]
                  |
                  v
[DataLoader: collate samples into a batch]
                  |
                  v
[Training loop: move batch to model device]
                  |
                  v
[albumentations.batch: MixUp or input normalization]
                  |
                  v
[Model inputs and targets for the loss]

Call batch operations directly in your training loop. Compose works earlier in the pipeline, where it can also synchronize transforms across collections such as video frames passed as images. A training batch contains separate examples that MixUp can pair with each other. Working with Targets

Batch operations run outside autograd on the batch device. Their outputs can feed a model that trains with gradients. This device contract is separate from Compose, which accepts NumPy or CPU Tensors; accepting FP16 in Compose does not make its transforms run on the GPU.

Mix examples and labels with one plan

MixUp blends each input with a donor from the same batch. The plan records a donor index and an original-sample weight for each example. Reuse it for both inputs and targets so they blend in the same proportions. The example converts class IDs to floating-point one-hot targets before mixing.

import torch
from torch.nn.functional import one_hot

from albumentations.batch import apply_mixup, sample_mixup

# Example output of collation. Use the tensors from your loader in training.
images = torch.arange(8, dtype=torch.float32).reshape(2, 1, 2, 2)
class_ids = torch.tensor([0, 1])

# Create this once per training loop on the same device as the batch.
generator = torch.Generator(device=images.device).manual_seed(137)
plan = sample_mixup(
    batch_size=images.shape[0],
    alpha=0.4,
    device=images.device,
    generator=generator,
)
mixed_images = apply_mixup(images, plan)
mixed_targets = apply_mixup(one_hot(class_ids, num_classes=2).float(), plan)

Use a loss that accepts class probabilities, such as PyTorch cross-entropy with floating-point targets shaped like the logits. For regression or segmentation, apply the plan to floating targets when interpolation is valid for the task. Detection boxes and keypoints require task-specific handling.

Mode	How donors and weights are chosen
`pair` (default)	Randomly pair examples and draw a separate weight for each example.
`element`	Draw a permutation of donors and a separate weight for each example.
`batch`	Randomly pair examples and use one weight for all pairs.

An example assigned to itself keeps weight 1. In pair and batch mode, an odd-sized batch leaves one example unpaired. Keep the generator across batches so its random stream advances. Apply MixUp during training.

apply_mixup uses the leading batch axis for both images (BCHW) and volumes (BCDHW). It interpolates in FP32 and returns the input dtype. For uint8 images, it rounds and clips the blended values to the byte range.

Normalize each sample on the model device

InputNormalizer takes channel-first image batches (BCHW) or volume batches (BCDHW). The modes determine which statistics it uses.

Mode	Operation
`fixed`	Divide by `input_scale`, subtract a supplied mean, and divide by the supplied standard deviation. Use scalar or per-channel statistics in the scaled intensity units.
`per_sample`	Compute statistics over the spatial axes of each sample and channel.
`per_sample_nonzero`	Use the nonzero intensities of each sample and channel. Preserve zeros in the output.
`none`	Scale and cast the input.

The output precision can be fp32, fp16, or bf16. Statistics and normalization arithmetic run in FP32 before the output cast. For fixed mode, move the module's statistics to the input device with normalizer.to(device=inputs.device).

For a zero-padded 3D crop, per_sample_nonzero keeps padding out of the mean and variance. Here the two measured intensities become −1 and 1, while the padded voxels stay zero.

import torch

from albumentations.batch import InputNormalizer

volumes = torch.tensor([0, 0, 10, 20], dtype=torch.uint8).reshape(1, 1, 1, 2, 2)
normalizer = InputNormalizer("per_sample_nonzero", output_precision="fp32")
normalized = normalizer(volumes)
print(normalized.flatten().tolist())  # [0.0, 0.0, -1.0, 1.0]