wehub-resource-sync
413 lines
16 KiB
Python
413 lines
16 KiB
Python
import sys
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray import train
|
|
from ray.data.preprocessors import Concatenator
|
|
from ray.train import ScalingConfig
|
|
|
|
if sys.version_info <= (3, 12):
|
|
# Skip this test for Python 3.12+ due to tensorflow incompatibility
|
|
import tensorflow as tf
|
|
|
|
# if tf version is > 2.16, errors cannot be imported as functions
|
|
# parse version with packaging
|
|
from packaging import version
|
|
|
|
from ray.train.tensorflow import TensorflowTrainer
|
|
|
|
if version.parse(tf.__version__) >= version.parse("2.16"):
|
|
mse = tf.keras.losses.MeanSquaredError()
|
|
mae = tf.keras.losses.MeanAbsoluteError()
|
|
else:
|
|
mse = tf.keras.losses.mean_squared_error
|
|
mae = tf.keras.losses.mean_absolute_error
|
|
|
|
|
|
class TestToTF:
|
|
def test_autosharding_is_disabled(self):
|
|
ds = ray.data.from_items([{"spam": 0, "ham": 0}])
|
|
|
|
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
|
|
|
actual_auto_shard_policy = (
|
|
dataset.options().experimental_distribute.auto_shard_policy
|
|
)
|
|
expected_auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
|
|
assert actual_auto_shard_policy is expected_auto_shard_policy
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_type(self, include_additional_columns):
|
|
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam", label_columns="ham", additional_columns="weight"
|
|
)
|
|
feature_spec, label_spec, additional_spec = dataset.element_spec
|
|
else:
|
|
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
|
feature_spec, label_spec = dataset.element_spec
|
|
|
|
assert isinstance(feature_spec, tf.TypeSpec)
|
|
assert isinstance(label_spec, tf.TypeSpec)
|
|
if include_additional_columns:
|
|
assert isinstance(additional_spec, tf.TypeSpec)
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_user_provided(self, include_additional_columns):
|
|
ds = ray.data.from_items([{"spam": 0, "ham": 0, "eggs": 0, "weight": 0}])
|
|
|
|
if include_additional_columns:
|
|
dataset1 = ds.to_tf(
|
|
feature_columns=["spam", "ham"],
|
|
label_columns="eggs",
|
|
additional_columns="weight",
|
|
)
|
|
feature_spec, label_spec, additional_spec = dataset1.element_spec
|
|
dataset2 = ds.to_tf(
|
|
feature_columns=["spam", "ham"],
|
|
label_columns="eggs",
|
|
additional_columns="weight",
|
|
feature_type_spec=feature_spec,
|
|
label_type_spec=label_spec,
|
|
additional_type_spec=additional_spec,
|
|
)
|
|
(
|
|
feature_output_spec,
|
|
label_output_spec,
|
|
additional_output_spec,
|
|
) = dataset2.element_spec
|
|
else:
|
|
dataset1 = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
|
|
feature_spec, label_spec = dataset1.element_spec
|
|
dataset2 = ds.to_tf(
|
|
feature_columns=["spam", "ham"],
|
|
label_columns="eggs",
|
|
feature_type_spec=feature_spec,
|
|
label_type_spec=label_spec,
|
|
)
|
|
feature_output_spec, label_output_spec = dataset2.element_spec
|
|
|
|
assert isinstance(label_output_spec, tf.TypeSpec)
|
|
assert isinstance(feature_output_spec, dict)
|
|
assert feature_output_spec.keys() == {"spam", "ham"}
|
|
assert all(
|
|
isinstance(value, tf.TypeSpec) for value in feature_output_spec.values()
|
|
)
|
|
if include_additional_columns:
|
|
assert isinstance(additional_output_spec, tf.TypeSpec)
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_type_with_multiple_columns(self, include_additional_columns):
|
|
ds = ray.data.from_items(
|
|
[{"spam": 0, "ham": 0, "eggs": 0, "weight1": 0, "weight2": 0}]
|
|
)
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns=["spam", "ham"],
|
|
label_columns="eggs",
|
|
additional_columns=["weight1", "weight2"],
|
|
)
|
|
(
|
|
feature_output_signature,
|
|
_,
|
|
additional_output_signature,
|
|
) = dataset.element_spec
|
|
else:
|
|
dataset = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
|
|
feature_output_signature, _ = dataset.element_spec
|
|
|
|
assert isinstance(feature_output_signature, dict)
|
|
assert feature_output_signature.keys() == {"spam", "ham"}
|
|
assert all(
|
|
isinstance(value, tf.TypeSpec)
|
|
for value in feature_output_signature.values()
|
|
)
|
|
|
|
if include_additional_columns:
|
|
assert isinstance(additional_output_signature, dict)
|
|
assert additional_output_signature.keys() == {"weight1", "weight2"}
|
|
assert all(
|
|
isinstance(value, tf.TypeSpec)
|
|
for value in additional_output_signature.values()
|
|
)
|
|
|
|
df = pd.DataFrame(
|
|
{
|
|
"feature1": [0, 1, 2],
|
|
"feature2": [3, 4, 5],
|
|
"label": [0, 1, 1],
|
|
"weight1": [0, 0.1, 0.2],
|
|
"weight2": [0.3, 0.4, 0.5],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(df)
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns=["feature1", "feature2"],
|
|
label_columns="label",
|
|
additional_columns=["weight1", "weight2"],
|
|
batch_size=3,
|
|
)
|
|
(
|
|
feature_output_signature,
|
|
_,
|
|
additional_output_signature,
|
|
) = dataset.element_spec
|
|
assert isinstance(additional_output_signature, dict)
|
|
assert additional_output_signature.keys() == {"weight1", "weight2"}
|
|
assert all(
|
|
isinstance(value, tf.TypeSpec)
|
|
for value in additional_output_signature.values()
|
|
)
|
|
else:
|
|
dataset = ds.to_tf(
|
|
feature_columns=["feature1", "feature2"],
|
|
label_columns="label",
|
|
batch_size=3,
|
|
)
|
|
feature_output_signature, _ = dataset.element_spec
|
|
|
|
assert isinstance(feature_output_signature, dict)
|
|
assert feature_output_signature.keys() == {"feature1", "feature2"}
|
|
assert all(
|
|
isinstance(value, tf.TypeSpec)
|
|
for value in feature_output_signature.values()
|
|
)
|
|
|
|
if include_additional_columns:
|
|
features, labels, additional_metadata = next(iter(dataset))
|
|
assert (
|
|
additional_metadata["weight1"].numpy() == df["weight1"].values
|
|
).all()
|
|
assert (
|
|
additional_metadata["weight2"].numpy() == df["weight2"].values
|
|
).all()
|
|
else:
|
|
features, labels = next(iter(dataset))
|
|
assert (labels.numpy() == df["label"].values).all()
|
|
assert (features["feature1"].numpy() == df["feature1"].values).all()
|
|
assert (features["feature2"].numpy() == df["feature2"].values).all()
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_name(self, include_additional_columns):
|
|
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam", label_columns="ham", additional_columns="weight"
|
|
)
|
|
feature_spec, label_spec, additional_spec = dataset.element_spec
|
|
else:
|
|
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
|
feature_spec, label_spec = dataset.element_spec
|
|
|
|
assert feature_spec.name == "spam"
|
|
assert label_spec.name == "ham"
|
|
if include_additional_columns:
|
|
assert additional_spec.name == "weight"
|
|
|
|
@pytest.mark.parametrize(
|
|
"data, expected_dtype",
|
|
# Skip this test for Python 3.12+ due to tensorflow incompatibility
|
|
[
|
|
(0, tf.int64),
|
|
(0.0, tf.double),
|
|
(False, tf.bool),
|
|
("eggs", tf.string),
|
|
([1.0, 2.0], tf.float64),
|
|
(np.zeros([2, 2], dtype=np.float32), tf.float32),
|
|
]
|
|
if sys.version_info <= (3, 12)
|
|
else [],
|
|
)
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_dtype(self, data, expected_dtype, include_additional_columns):
|
|
ds = ray.data.from_items([{"spam": data, "ham": data, "weight": data}])
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam",
|
|
label_columns="ham",
|
|
additional_columns="weight",
|
|
)
|
|
feature_spec, label_spec, additional_spec = dataset.element_spec
|
|
else:
|
|
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
|
feature_spec, label_spec = dataset.element_spec
|
|
|
|
assert feature_spec.dtype == expected_dtype
|
|
assert label_spec.dtype == expected_dtype
|
|
if include_additional_columns:
|
|
assert additional_spec.dtype == expected_dtype
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_shape(self, include_additional_columns):
|
|
ds = ray.data.from_items(8 * [{"spam": 0, "ham": 0, "weight": 0}])
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam",
|
|
label_columns="ham",
|
|
additional_columns="weight",
|
|
batch_size=4,
|
|
)
|
|
feature_spec, label_spec, additional_spec = dataset.element_spec
|
|
assert tuple(additional_spec.shape) == (None,)
|
|
else:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam", label_columns="ham", batch_size=4
|
|
)
|
|
feature_spec, label_spec = dataset.element_spec
|
|
|
|
assert tuple(feature_spec.shape) == (None,)
|
|
assert tuple(label_spec.shape) == (None,)
|
|
|
|
if include_additional_columns:
|
|
features, labels, additional_metadata = next(iter(dataset))
|
|
assert tuple(additional_metadata.shape) == (4,)
|
|
else:
|
|
features, labels = next(iter(dataset))
|
|
assert tuple(features.shape) == (4,)
|
|
assert tuple(labels.shape) == (4,)
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_shape_with_tensors(self, include_additional_columns):
|
|
ds = ray.data.from_items(
|
|
8
|
|
* [
|
|
{
|
|
"spam": np.zeros([3, 32, 32]),
|
|
"ham": 0,
|
|
"weight": np.zeros([3, 32, 32]),
|
|
}
|
|
]
|
|
)
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam",
|
|
label_columns="ham",
|
|
additional_columns="weight",
|
|
batch_size=4,
|
|
)
|
|
feature_spec, _, additional_spec = dataset.element_spec
|
|
assert tuple(additional_spec.shape) == (None, 3, 32, 32)
|
|
else:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam", label_columns="ham", batch_size=4
|
|
)
|
|
feature_spec, _ = dataset.element_spec
|
|
|
|
assert tuple(feature_spec.shape) == (None, 3, 32, 32)
|
|
|
|
if include_additional_columns:
|
|
features, labels, additional_metadata = next(iter(dataset))
|
|
assert tuple(additional_metadata.shape) == (4, 3, 32, 32)
|
|
else:
|
|
features, labels = next(iter(dataset))
|
|
assert tuple(features.shape) == (4, 3, 32, 32)
|
|
assert tuple(labels.shape) == (4,)
|
|
|
|
@pytest.mark.parametrize("batch_size", [1, 2])
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_element_spec_shape_with_ragged_tensors(
|
|
self, batch_size, include_additional_columns
|
|
):
|
|
df = pd.DataFrame(
|
|
{
|
|
"spam": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
|
|
"ham": [0, 0],
|
|
"weight": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(df)
|
|
|
|
if include_additional_columns:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam",
|
|
label_columns="ham",
|
|
additional_columns="weight",
|
|
batch_size=batch_size,
|
|
)
|
|
feature_spec, _, additional_spec = dataset.element_spec
|
|
assert tuple(additional_spec.shape) == (None, None, None, None)
|
|
else:
|
|
dataset = ds.to_tf(
|
|
feature_columns="spam", label_columns="ham", batch_size=batch_size
|
|
)
|
|
feature_spec, _ = dataset.element_spec
|
|
|
|
assert tuple(feature_spec.shape) == (None, None, None, None)
|
|
|
|
if include_additional_columns:
|
|
features, labels, additional_metadata = next(iter(dataset))
|
|
assert tuple(additional_metadata.shape) == (batch_size, None, None, None)
|
|
else:
|
|
features, labels = next(iter(dataset))
|
|
assert tuple(features.shape) == (batch_size, None, None, None)
|
|
assert tuple(labels.shape) == (batch_size,)
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_training(self, include_additional_columns):
|
|
def build_model() -> tf.keras.Model:
|
|
return tf.keras.Sequential([tf.keras.layers.Dense(1)])
|
|
|
|
def train_func():
|
|
strategy = tf.distribute.MultiWorkerMirroredStrategy()
|
|
with strategy.scope():
|
|
multi_worker_model = build_model()
|
|
multi_worker_model.compile(
|
|
optimizer=tf.keras.optimizers.SGD(),
|
|
loss=mae,
|
|
metrics=[mse],
|
|
)
|
|
|
|
if include_additional_columns:
|
|
dataset = train.get_dataset_shard("train").to_tf(
|
|
"X", "Y", additional_columns="W", batch_size=4
|
|
)
|
|
else:
|
|
dataset = train.get_dataset_shard("train").to_tf("X", "Y", batch_size=4)
|
|
multi_worker_model.fit(dataset)
|
|
|
|
dataset = ray.data.from_items(8 * [{"X0": 0, "X1": 0, "Y": 0, "W": 0}])
|
|
concatenator = Concatenator(columns=["X0", "X1"], output_column_name="X")
|
|
dataset = concatenator.transform(dataset)
|
|
|
|
trainer = TensorflowTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(num_workers=2),
|
|
datasets={"train": dataset},
|
|
)
|
|
trainer.fit()
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_invalid_column_raises_error(self, include_additional_columns):
|
|
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
|
with pytest.raises(ValueError):
|
|
if include_additional_columns:
|
|
ds.to_tf(
|
|
feature_columns="spam",
|
|
label_columns="ham",
|
|
additional_columns="baz",
|
|
)
|
|
else:
|
|
ds.to_tf(feature_columns="foo", label_columns="bar")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if sys.version_info >= (3, 12):
|
|
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
|
sys.exit(0)
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|