Release 2.6.2

Summary

  • Augment dense gradients, displacements, local vectors, and surface normals alongside their source images or volumes.
  • Keep field components consistent with the sampled geometry and preserve invalid regions with NaN-aware interpolation.
  • Reduce latency and temporary GPU allocation for eligible FP16 and BF16 batch normalization on Linux CUDA.

Breaking changes

Vector-field target names

gradient, displacement, vector_field, normal, and their 3d counterparts are now built-in spatial targets. If a permissive pipeline used one of these undeclared keys for unrelated metadata, move the value into user_data and read it from the returned user_data. Strict pipelines should omit absent field targets. #635

# Before
result = pipeline(image=image, gradient=metadata)

# After
result = pipeline(image=image, user_data={"gradient": metadata})
metadata = result["user_data"]["gradient"]

New features

Geometry-aware vector fields

Pass fields to the same Compose call as the image or volume. Each supported transform uses the same sampled geometry for the raster and its fields, including when replaying a pipeline. Choose the target that describes the field's values.

2D target3D targetComponent update
gradientgradient3dApply the inverse transpose of the local Jacobian
displacementdisplacement3dMap the start and end points and subtract their new coordinates
vector_fieldvector_field3dApply the local Jacobian
normalnormal3dApply the inverse transpose and restore unit length

Doubling both image dimensions doubles a displacement measured in pixels and halves a scalar-field gradient. This example resizes optical flow with its image. Vector-field guide · #635

import albumentations as A
import numpy as np

image = np.zeros((32, 48, 3), dtype=np.uint8)
flow = np.full((32, 48, 2), (2.0, 1.0), dtype=np.float32)
augment = A.Compose([A.Resize(height=64, width=96)])
result = augment(image=image, displacement=flow)

assert result["displacement"].shape == (64, 96, 2)
np.testing.assert_allclose(result["displacement"][32, 48], (4.0, 2.0))
  • Fields accept float16 and float32. NumPy uses a trailing component axis and CPU PyTorch tensors use a leading component axis. A 2D normal has three components on a 2D grid. Use additional_targets for extra named fields. ToTensorV2 and ToTensor3D preserve field dtype and move components to the leading axis. Field values can be negative or greater than one.
  • Resampling uses bilinear or trilinear interpolation and renormalizes weights over finite neighbors. Output outside the input grid and positions with no valid support become NaN. Padding adds NaNs to fields.
  • Dropout preserves fields. Check each transform's Targets section before adding fields to an existing pipeline. The following transforms skip the entire node when a field is present. CopyAndPaste · Mosaic · OverlayElements · Morphological · PixelSpread

Performance

Fused CUDA normalization

InputNormalizer automatically uses the fused path for eligible tensors on Linux when Triton is available and the NVIDIA GPU has compute capability 8.0 or later. The kernels normalize in FP32, write the final reduced-precision output directly, and preserve the input memory format.

ModeFused input and output combinations
fixeduint8 or FP16 → FP16 or BF16. FP32 → BF16.
per_sample, per_sample_nonzerouint8 → FP16 or BF16. FP16 → BF16.

Calls inside torch.compile, ineligible layouts or sizes, and recognized Triton compilation failures use the PyTorch path. Both z-score modes use tiled FP32 statistics on the fused path, so compare their results with numerical tolerances. CUDA normalization contract · #643

On an NVIDIA L4 with PyTorch 2.15.0+cu132 and Triton 3.9.0, the PR measured these resident GPU batches against eager AlbumentationsX at 7216f53. Peak allocation is measured above the resident input. Some small image batches remained faster in eager PyTorch, particularly with standard z-score. Full benchmark results

Batch and modeEager → fused medianEager → fused peak allocation
2 × 7 × 384 × 160 × 160, fixed FP32 → BF1617.50 → 3.62 ms788.5 → 262.5 MiB
288 × 1 × 32 × 384 × 384, nonzero z-score FP16 → BF16297.45 → 34.46 ms7,776 → 2,599.6 MiB

Misc

  • Updated the development Ruff hook from 0.17.3 to 0.17.4 and the Pyrefly hook from 1.4.0 to 1.4.1. #635

Commits

CommitPRDescription
7216f53#635Add 2D and 3D vector-field targets
76f0735#643Fuse CUDA batch normalization for FP16 and BF16 outputs
0237b80#644Bump to 2.6.2 and update the migration inventory

Full changelog for 2.6.1 to 2.6.2