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 target | 3D target | Component update |
|---|---|---|
gradient | gradient3d | Apply the inverse transpose of the local Jacobian |
displacement | displacement3d | Map the start and end points and subtract their new coordinates |
vector_field | vector_field3d | Apply the local Jacobian |
normal | normal3d | Apply 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
float16andfloat32. 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. Useadditional_targetsfor extra named fields.ToTensorV2andToTensor3Dpreserve 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
Targetssection 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.
| Mode | Fused input and output combinations |
|---|---|
fixed | uint8 or FP16 → FP16 or BF16. FP32 → BF16. |
per_sample, per_sample_nonzero | uint8 → 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 mode | Eager → fused median | Eager → fused peak allocation |
|---|---|---|
2 × 7 × 384 × 160 × 160, fixed FP32 → BF16 | 17.50 → 3.62 ms | 788.5 → 262.5 MiB |
288 × 1 × 32 × 384 × 384, nonzero z-score FP16 → BF16 | 297.45 → 34.46 ms | 7,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