Is this page helpful?

Implementing custom operators#

A pre-processing pipeline may need a transformation that isn’t available in TorchVision (and is therefore also missing from the DALI TorchVision API). In that case, you can implement a custom operator. Implemented correctly, it can be reused across pipelines just like a built-in one, by adding it to a Compose list.

Warning

The extension mechanism described here is experimental.

Operator, ArgumentValidateRule, adjust_input and the layout helpers are imported from the nvidia.dali.experimental.torchvision.v2.operator module, and the TensorLike and DeviceLike type aliases come from internal DALI modules. None of them is exported from the public nvidia.dali.experimental.torchvision package, so they may change or be moved in a future release without notice.

[1]:

import numpy as np
import torch
from PIL import Image

import nvidia.dali as dali
import nvidia.dali.fn as fn
import nvidia.dali.experimental.dynamic as ndd
import nvidia.dali.experimental.torchvision as dtv
from nvidia.dali._typing import TensorLike
from nvidia.dali.experimental.dynamic._device import DeviceLike
from nvidia.dali.experimental.torchvision.v2.operator import (
    ArgumentValidateRule,
    Operator,
    adjust_input,
    get_HWC_from_layout_dynamic,
    get_HWC_from_layout_pipeline,
)


def _stack_int64(*values):
    """Helper: pack scalar DALI values into a single int64 vector."""
    return fn.stack(*[fn.cast(v, dtype=dali.types.INT64) for v in values])

Object-oriented operators#

Processing flow#

To implement a custom operator, derive from the Operator class and implement at least the _kernel method, which contains the core processing logic. There are also three optional hooks:

  • arg_rules — validates the transformation’s configuration parameters when the operator is created.

  • input_rules — validates the input data at the transformation’s runtime.

  • preprocess_data — pre-processes the input data before it is passed to the core pipeline logic.

Both arg_rules and input_rules are lists of callable objects (ArgumentValidateRule and DataValidateRule subclasses, respectively) that check whether the data adheres to a rule. See Validating arguments and input for details.

Keep in mind that preprocess_data may change the form of the original input data, so _kernel may need to adjust to it. For example, get_HWC_from_layout_pipeline outputs a tuple (H, W, C, input_data_tensor) instead of just the tensor, as the TrimBorder example below does.

When you call an instance of Compose, the data is passed through the transformations from the first to the last on the list. Each transformation:

  1. validates the input data with input_rules,

  2. copies the input data to the GPU if the transformation’s device is "gpu",

  3. runs preprocess_data,

  4. passes the result to _kernel, and its output becomes the input of the next transformation.

The _invoke method of the Operator class, shown further below, implements this flow.

Example#

An example custom operator is shown below. It removes a border of border_size pixels from each side of the image. First, a rule that validates its argument:

[2]:

class ValidateBorderSize(ArgumentValidateRule):
    @classmethod
    def verify(cls, *, border_size, **_):
        if not isinstance(border_size, int) or border_size < 0:
            raise ValueError(
                f"border_size must be a non-negative integer, got {border_size!r}."
            )

[3]:

class TrimBorder(Operator):
    arg_rules = [ValidateBorderSize]
    preprocess_data = get_HWC_from_layout_pipeline

    def __init__(
        self,
        border_size: int,
        device: str = "cpu",
    ):
        super().__init__(device=device, border_size=border_size)
        self.border_size = border_size

    def _kernel(self, data_input):
        in_h, in_w, _, tensor = data_input
        top = dali.math.min(self.border_size, in_h)
        left = dali.math.min(self.border_size, in_w)
        height = dali.math.max(0, in_h - 2 * self.border_size)
        width = dali.math.max(0, in_w - 2 * self.border_size)
        return fn.slice(
            tensor,
            _stack_int64(left, top),
            _stack_int64(width, height),
            axis_names="WH",
            out_of_bounds_policy="pad",
            fill_values=0,
            device=self.device,
        )

TrimBorder is a subclass of Operator and implements the two required methods: the constructor __init__ and _kernel. The latter contains the DALI code that implements the transformation logic. When listed as a step in Compose, it becomes part of the pipeline just like any built-in operator, as shown below:

[4]:

# A PIL image: the pipeline works on HWC data and returns a PIL image
pil_pipeline = dtv.Compose([TrimBorder(border_size=4)])
pil_img = Image.fromarray(np.zeros((20, 30, 3), dtype=np.uint8))
print(pil_pipeline(pil_img).size)  # (22, 12)

# A torch.Tensor: the pipeline works on CHW data and returns a torch.Tensor
tensor_pipeline = dtv.Compose([TrimBorder(border_size=4)])
print(tensor_pipeline(torch.rand(3, 20, 30)).shape)  # torch.Size([3, 12, 22])

# Invalid arguments are rejected when the operator is created
try:
    TrimBorder(border_size=-1)
except ValueError as e:
    print(e)  # border_size must be a non-negative integer, got -1.

(22, 12)
torch.Size([3, 12, 22])
border_size must be a non-negative integer, got -1.

Operator is the base class for all object-oriented transformations. The following code highlights its main design points:

class Operator(ABC):
    """
    Abstract base class for operator specification

    Implement _kernel for algorithm specific processing

    ``arg_rules`` - a sequence of verification rules for algorithm's arguments.
    ``input_rules`` - a sequence of verification rules for algorithm's input data.
    ``preprocess_data`` - a function to preprocess the input data.

    Parameters
    ----------
    device : Literal["cpu", "gpu"], optional, default = "cpu"
        Device to use for the operator. Can be ``"cpu"`` or ``"gpu"``.
    **kwargs
        Additional keyword arguments for the operator.
    """

    arg_rules: tuple[ArgumentValidateRule, ...] = tuple()
    input_rules: tuple[DataValidateRule, ...] = tuple()
    preprocess_data = None

    @classmethod
    def verify_args(cls, **kwargs):
        for rule in cls.arg_rules:
            rule.verify(**kwargs)

    @classmethod
    def verify_data(cls, data_input):
        for rule in cls.input_rules:
            rule.verify(data_input)

    def __init__(self, device: Literal["cpu", "gpu"] = "cpu", **kwargs):
        self.device = device
        type(self).verify_args(**kwargs)

    @abstractmethod
    def _kernel(self, data_input):
        """
        Algorithm's processing
        """
        pass

    ...

    def _invoke(self, data_input):
        """
        Private method is used to execute operator from a Compose pipeline
        Note: Do not call directly
        """
        type(self).verify_data(data_input)

        # Original input is transferred to GPU, before being preprocess_data.
        # The preprocess_data creates an arbitrary tuple
        if self.device == "gpu":
            data_input = data_input.gpu()

        if type(self).preprocess_data:
            data_input = type(self).preprocess_data(data_input)

        output = self._kernel(data_input)

        return output

Validating arguments and input#

A minimal custom operator only needs to implement _kernel, which contains the transformation’s logic. If the operator takes configuration parameters, implement __init__ as well. The remaining class attributes are optional:

  • arg_rules — rules that validate the operator’s arguments, as TrimBorder does with ValidateBorderSize. Each rule derives from ArgumentValidateRule and implements verify as a class method that raises an exception when the check fails. verify receives the keyword arguments that the operator passes to super().__init__. Every rule in arg_rules receives all of them, so a rule should name the arguments it checks and accept the rest with **_. Because the rules run in the constructor, a wrong argument fails early, when the operator is created, and not when the pipeline is executed. Existing validators can be reused in the same way.

  • input_rules — rules that validate the input data. Each rule derives from DataValidateRule and implements verify(cls, data) as a class method. Inside a Compose pipeline, data is a DALI DataNode that describes the pipeline graph, not a torch.Tensor or a PIL.Image, so such a rule can’t inspect the actual shape or values. Check the arguments with arg_rules instead, and handle input-dependent cases in _kernel.

  • preprocess_data — a function that processes the input data before it is passed to _kernel. TrimBorder uses get_HWC_from_layout_pipeline, which makes _kernel receive the height, width and number of channels of the input together with the data itself.

An important limitation is that arg_rules can’t see the input image. TrimBorder doesn’t reject a border_size that is at least half of the image height or width. In that case, the result has zero height or width, which is the same as for the functional version below. When in doubt, check how the existing operators are implemented.

Functional operators#

The functional API is simpler, with no formal rules for implementing a custom operator. However, when DALI Dynamic Mode is used for processing, operators are expected to consume a DALI Tensor or Batch. Input data can be adjusted automatically by decorating the processing function with the @adjust_input decorator, as shown in the example below. It implements the same TrimBorder transformation as the operator above and reuses the ValidateBorderSize rule to check its argument:

[5]:

@adjust_input
def trim_border(
    inpt: TensorLike | ndd.Batch,
    border_size: int,
    device: DeviceLike = "cpu",
) -> ndd.Tensor | ndd.Batch:
    ValidateBorderSize.verify(border_size=border_size)
    in_h, in_w, _ = get_HWC_from_layout_dynamic(inpt)
    top = min(border_size, in_h)
    left = min(border_size, in_w)
    height = max(0, in_h - 2 * border_size)
    width = max(0, in_w - 2 * border_size)
    axes = [-3, -2] if inpt.layout[-3:] == "HWC" else [-2, -1]
    return ndd.slice(
        inpt,
        (top, left),
        (height, width),
        axes=axes,
        out_of_bounds_policy="pad",
        fill_values=0,
        device=device,
    )

[6]:

pil_img = Image.fromarray(np.zeros((20, 30, 3), dtype=np.uint8))
print(trim_border(pil_img, border_size=4).size)  # (22, 12)

print(
    trim_border(torch.rand(3, 20, 30), border_size=4).shape
)  # torch.Size([3, 12, 22])

# A batch of images: a tensor with a leading batch dimension (N, C, H, W)
print(
    trim_border(torch.rand(2, 3, 20, 30), border_size=4).shape
)  # torch.Size([2, 3, 12, 22])

(22, 12)
torch.Size([3, 12, 22])
torch.Size([2, 3, 12, 22])