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.

RGB input pipeline results

Every path reads RGB JPEGs, prepares the recipe, and delivers a synchronized CUDA batch. CPU and GPU labels identify where augmentation runs; normalization runs on GPU for every path.

Each comparison uses its own shared recipe set. Averages from different sets cannot rank all libraries. The table below includes every measured recipe for these paths, including recipes outside the summary set. Recipe names are shortened; hover over a name for its full pipeline.

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"]