Working with tensors and NumPy#

Tensors, or n-dimensional arrays, are ubiquitous in machine learning workloads. This guide describes how Ray Data represents tensor data and best practices for working with it.

How does Ray Data represent tensors?#

Ray Data represents tensors as NumPy ndarrays.

import ray

ds = ray.data.read_images("s3://anonymous@air-example-data/digits")
print(ds)

Dataset(num_rows=100, schema=...)

How does Ray Data batch fixed-shape tensors?#

If your tensors have a fixed shape, Ray Data represents batches as regular ndarrays.

>>> import ray
>>> ds = ray.data.read_images("s3://anonymous@air-example-data/digits")
>>> batch = ds.take_batch(batch_size=32)
>>> batch["image"].shape
(32, 28, 28)
>>> batch["image"].dtype
dtype('uint8')

How does Ray Data batch variable-shape tensors?#

If your tensors vary in shape, Ray Data represents batches as arrays of object dtype.

>>> import ray
>>> ds = ray.data.read_images("s3://anonymous@air-example-data/AnimalDetection")
>>> batch = ds.take_batch(batch_size=32)
>>> batch["image"].shape
(32,)
>>> batch["image"].dtype
dtype('O')

Each element of these object arrays is a regular ndarray.

>>> batch["image"][0].dtype
dtype('uint8')
>>> batch["image"][0].shape
(375, 500, 3)
>>> batch["image"][3].shape
(333, 465, 3)

Transform tensor data#

Call map() or map_batches() to transform tensor data.

from typing import Any, Dict

import ray
import numpy as np

ds = ray.data.read_images("s3://anonymous@air-example-data/AnimalDetection")

def increase_brightness(row: Dict[str, Any]) -> Dict[str, Any]:
    row["image"] = np.clip(row["image"] + 4, 0, 255)
    return row

# Increase the brightness, record at a time.
ds.map(increase_brightness)

def batch_increase_brightness(batch: Dict[str, np.ndarray]) -> Dict:
    batch["image"] = np.clip(batch["image"] + 4, 0, 255)
    return batch

# Increase the brightness, batch at a time.
ds.map_batches(batch_increase_brightness, batch_size="auto")

Set batch_size="auto" to have Ray Data pick a batch size based on the size of your data. When you set num_gpus, batch_size must be an integer instead.

Besides NumPy ndarrays, Ray Data also treats returned lists of NumPy ndarrays as tensor data. The same goes for returned objects that implement __array__, such as torch.Tensor.

For more information on transforming data, see Transforming data.

Save tensor data#

Save tensor data in formats such as Parquet, NumPy, and JSON. For all supported formats, see the Saving Data API.

Call write_parquet() to save data in Parquet files.

import ray

ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
ds.write_parquet("/tmp/simple")

Call write_numpy() to save an ndarray column in NumPy files.

import ray

ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
ds.write_numpy("/tmp/simple", column="image")

Call write_json() to save images in a JSON file.

import ray

ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
ds.write_json("/tmp/simple")

For more information on saving data, see Saving data.