ray.data.preprocessors.TorchVisionPreprocessor#

class ray.data.preprocessors.TorchVisionPreprocessor(columns: List[str], transform: Callable[[np.ndarray | torch.Tensor], torch.Tensor], output_columns: List[str] | None = None, batched: bool = False)[source]#

Bases: SerializablePreprocessorBase

Apply a TorchVision transform to image columns.

Examples

Torch models expect inputs of shape \((B, C, H, W)\) in the range \([0.0, 1.0]\). To convert images to this format, add ToTensor to your preprocessing pipeline.

from torchvision import transforms

import ray
from ray.data.preprocessors import TorchVisionPreprocessor

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((224, 224)),
])
preprocessor = TorchVisionPreprocessor(["image"], transform=transform)

dataset = ray.data.read_images("s3://anonymous@air-example-data-2/imagenet-sample-images")
dataset = preprocessor.transform(dataset)

For better performance, set batched to True and replace ToTensor with a batch-supporting Lambda.

import numpy as np
import torch

def to_tensor(batch: np.ndarray) -> torch.Tensor:
    tensor = torch.as_tensor(batch, dtype=torch.float)
    # (B, H, W, C) -> (B, C, H, W)
    tensor = tensor.permute(0, 3, 1, 2).contiguous()
    # [0., 255.] -> [0., 1.]
    tensor = tensor.div(255)
    return tensor

transform = transforms.Compose([
    transforms.Lambda(to_tensor),
    transforms.Resize((224, 224))
])
preprocessor = TorchVisionPreprocessor(["image"], transform=transform, batched=True)

dataset = ray.data.read_images("s3://anonymous@air-example-data-2/imagenet-sample-images")
dataset = preprocessor.transform(dataset)

Parameters:

columns – The columns to apply the TorchVision transform to.
transform – The TorchVision transform you want to apply. This transform should accept a np.ndarray or torch.Tensor as input and return a torch.Tensor as output.
output_columns – The output name for each input column. If not specified, this defaults to the same set of columns as the columns.
batched – If True, apply transform to batches of shape \((B, H, W, C)\). Otherwise, apply transform to individual images.

PublicAPI (alpha): This API is in alpha and may change before becoming stable.

Methods

`deserialize`	Deserialize a preprocessor from serialized data.
`fit`	Fit this Preprocessor to the Dataset.
`fit_transform`	Fit this Preprocessor to the Dataset and then transform the Dataset.
`get_preprocessor_class_id`	Get the preprocessor class identifier for this preprocessor class.
`get_version`	Get the version number for this preprocessor class.
`serialize`	Serialize this preprocessor to a string or bytes.
`set_preprocessor_class_id`	Set the preprocessor class identifier for this preprocessor class.
`set_version`	Set the version number for this preprocessor class.
`transform`	Transform the given dataset.
`transform_batch`	Transform a single batch of data.

Attributes

`MAGIC_CLOUDPICKLE`
`SERIALIZER_FORMAT_VERSION`
`batched`
`columns`
`output_columns`
`stat_computation_plan`
`torchvision_transform`