to_tf#

Return a TensorFlow Dataset over this Dataset.

Warning

If your Dataset contains ragged tensors, this method errors. To prevent errors, resize your tensors.

Note

This operation will trigger execution of the lazy transformations performed on this dataset.

Examples

>>> import ray
>>> ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
>>> ds
Dataset(num_rows=?, schema=...)

If your model accepts a single tensor as input, specify a single feature column.

>>> ds.to_tf(feature_columns="sepal length (cm)", label_columns="target")
<_OptionsDataset element_spec=(TensorSpec(shape=(None,), dtype=tf.float64, name='sepal length (cm)'), TensorSpec(shape=(None,), dtype=tf.int64, name='target'))>

If your model accepts a dictionary as input, specify a list of feature columns.

>>> ds.to_tf(["sepal length (cm)", "sepal width (cm)"], "target")
<_OptionsDataset element_spec=({'sepal length (cm)': TensorSpec(shape=(None,), dtype=tf.float64, name='sepal length (cm)'), 'sepal width (cm)': TensorSpec(shape=(None,), dtype=tf.float64, name='sepal width (cm)')}, TensorSpec(shape=(None,), dtype=tf.int64, name='target'))>

If your dataset contains multiple features but your model accepts a single tensor as input, combine features with Concatenator.

>>> from ray.data.preprocessors import Concatenator
>>> columns_to_concat = ["sepal length (cm)", "sepal width (cm)", "petal length (cm)", "petal width (cm)"]
>>> preprocessor = Concatenator(columns=columns_to_concat, output_column_name="features")
>>> ds = preprocessor.transform(ds)
>>> ds
Concatenator
+- Dataset(num_rows=?, schema=...)
>>> ds.to_tf("features", "target")
<_OptionsDataset element_spec=(TensorSpec(shape=(None, 4), dtype=tf.float64, name='features'), TensorSpec(shape=(None,), dtype=tf.int64, name='target'))>

If your model accepts different types, shapes, or names of tensors as input, specify the type spec. If type specs are not specified, they are automatically inferred from the schema of the dataset.

>>> import tensorflow as tf
>>> ds.to_tf(
...     feature_columns="features",
...     label_columns="target",
...     feature_type_spec=tf.TensorSpec(shape=(None, 4), dtype=tf.float32, name="features"),
...     label_type_spec=tf.TensorSpec(shape=(None,), dtype=tf.float32, name="label")
... )
<_OptionsDataset element_spec=(TensorSpec(shape=(None, 4), dtype=tf.float32, name='features'), TensorSpec(shape=(None,), dtype=tf.float32, name='label'))>

If your model accepts additional metadata aside from features and label, specify a single additional column or a list of additional columns. A common use case is to include sample weights in the data samples and train a tf.keras.Model with tf.keras.Model.fit.

>>> import pandas as pd
>>> ds = ds.add_column("sample weights", lambda df: pd.Series([1] * len(df)))
>>> ds.to_tf(feature_columns="features", label_columns="target", additional_columns="sample weights")
<_OptionsDataset element_spec=(TensorSpec(shape=(None, 4), dtype=tf.float64, name='features'), TensorSpec(shape=(None,), dtype=tf.int64, name='target'), TensorSpec(shape=(None,), dtype=tf.int64, name='sample weights'))>

If your model accepts different types, shapes, or names for the additional metadata, specify the type spec of the additional column.

>>> ds.to_tf(
...     feature_columns="features",
...     label_columns="target",
...     additional_columns="sample weights",
...     additional_type_spec=tf.TensorSpec(shape=(None,), dtype=tf.float32, name="weight")
... )
<_OptionsDataset element_spec=(TensorSpec(shape=(None, 4), dtype=tf.float64, name='features'), TensorSpec(shape=(None,), dtype=tf.int64, name='target'), TensorSpec(shape=(None,), dtype=tf.float32, name='weight'))>

Parameters:

feature_columns (str | List[str]) – Columns that correspond to model inputs. If this is a string, the input data is a tensor. If this is a list, the input data is a dict that maps column names to their tensor representation.
label_columns (str | List[str]) – Columns that correspond to model targets. If this is a string, the target data is a tensor. If this is a list, the target data is a dict that maps column names to their tensor representation.
additional_columns (str | List[str] | None) – Columns that correspond to sample weights or other metadata. If this is a string, the weight data is a tensor. If this is a list, the weight data is a dict that maps column names to their tensor representation.
prefetch_batches (int) – The number of batches to fetch ahead of the current batch. Ray Data prepares batches in background threads while your code consumes the current one. If greater than 0, a thread pool formats the batches. The pool has prefetch_batches threads, up to a limit of 4. To change the limit, set the RAY_DATA_MAX_FORMAT_THREADPOOL_NUM_WORKERS environment variable. If 0, Ray Data doesn’t prefetch blocks and formats each batch without the thread pool. Defaults to 1.
batch_size (int) – Record batch size. Defaults to 1.
drop_last (bool) – Set to True to drop the last incomplete batch, if the dataset size is not divisible by the batch size. If False and the size of the stream is not divisible by the batch size, then the last batch is smaller. Defaults to False.
local_shuffle_buffer_size (int | None) – If not None, Ray Data randomly shuffles the data using a local in-memory shuffle buffer, and this value serves as the minimum number of rows that must be in the buffer to yield a batch. When there are no more rows to add to the buffer, Ray Data drains the remaining rows. A larger buffer makes the shuffle more random but might increase the latency to the first batch. If this value is smaller than batch_size, Ray Data rounds it up to batch_size. Local shuffling raises a ValueError if batch_size is None.
local_shuffle_seed (int | None) – The seed to use for the local random shuffle.
feature_type_spec (tf.TypeSpec | Dict[str, tf.TypeSpec] | None) – The tf.TypeSpec of feature_columns. If there is only one column, specify a tf.TypeSpec. If there are multiple columns, specify a dict that maps column names to their tf.TypeSpec. Default is None to automatically infer the type of each column.
label_type_spec (tf.TypeSpec | Dict[str, tf.TypeSpec] | None) – The tf.TypeSpec of label_columns. If there is only one column, specify a tf.TypeSpec. If there are multiple columns, specify a dict that maps column names to their tf.TypeSpec. Default is None to automatically infer the type of each column.
additional_type_spec (tf.TypeSpec | Dict[str, tf.TypeSpec] | None) – The tf.TypeSpec of additional_columns. If there is only one column, specify a tf.TypeSpec. If there are multiple columns, specify a dict that maps column names to their tf.TypeSpec. Default is None to automatically infer the type of each column.

Returns:

A TensorFlow Dataset that yields inputs and targets.

Return type:

tf.data.Dataset