DALI TorchVision API

Starting from DALI 2.1, an experimental TorchVision-compatible API is available in the nvidia.dali.experimental.torchvision module. It provides GPU-accelerated drop-in replacements for torchvision.transforms.v2, so existing TorchVision preprocessing code can be offloaded to DALI with minimal changes.

For the full description, the list of supported operators and instructions on implementing custom operators, see the DALI TorchVision documentation.

Drop-in replacement

Switch the import - the rest of the code stays the same:

import torch

# import torchvision.transforms.v2 as T
import nvidia.dali.experimental.torchvision as T

preprocess = T.Compose(
    [
        T.Resize(size=256),
        T.CenterCrop(size=(224, 224)),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

image = torch.rand(3, 480, 640)  # CHW
output = preprocess(image)
print(output.shape)

torch.Size([3, 224, 224])

DALI adds two optional parameters: device ("cpu" or "gpu") for each operator and batch_size for Compose. Class-based operators must be wrapped in Compose.