AlbumentationsX vs Kornia

What Is Different?

Kornia is a differentiable computer vision library for PyTorch tensors. AlbumentationsX is a fast CPU augmentation library for NumPy arrays before data reaches the model.

• Kornia is tensor-first and shines on batched GPU augmentation; AlbumentationsX is NumPy-first and shines in CPU data-loading pipelines.
• Kornia transforms can be differentiable and participate in model graphs; AlbumentationsX transforms are preprocessing/data augmentation steps.
• AlbumentationsX has broad task-level target handling for images, masks, boxes, keypoints, volumes, and multiple related inputs.
• Kornia is a better fit when augmentation must happen on-device after batching; AlbumentationsX is usually simpler before batching.

Benchmark vs Kornia

The benchmark below compares single-threaded RGB image augmentation throughput. Kornia receives normalized tensors; AlbumentationsX receives OpenCV-loaded RGB NumPy arrays. Higher images/second is better.

See the aggregate image benchmark or inspect the benchmark source code.

Conversion Guide

The conversion usually moves augmentation from tensor batches in the training step into the Dataset/DataLoader preprocessing path.

• Apply AlbumentationsX before converting the sample to a tensor.
• Use ToTensorV2 at the end of the pipeline if the model expects PyTorch tensors.
• Move per-batch GPU-only transforms to AlbumentationsX only when they do not require differentiability or batched tensor semantics.
• Keep Kornia for model-integrated or differentiable computer vision operations.

import kornia.augmentation as K
import torch.nn as nn

augment = nn.Sequential(
    K.RandomHorizontalFlip(p=0.5),
    K.ColorJitter(brightness=0.2, contrast=0.2, p=0.5),
)

batch = augment(batch)

import albumentations as A
from albumentations.pytorch import ToTensorV2

transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
    ToTensorV2(),
])

sample = transform(image=image_np)
image = sample["image"]