Module keras.distribute.dataset_creator_model_fit_test_base
Tests for DatasetCreator
with Model.fit
across usages and strategies.
Expand source code
# Lint as: python3
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for `DatasetCreator` with `Model.fit` across usages and strategies."""
import tensorflow.compat.v2 as tf
import os
from absl.testing import parameterized
import numpy as np
import keras
from keras import callbacks as callbacks_lib
from keras.engine import sequential
from keras.layers import core as core_layers
from keras.layers.preprocessing import string_lookup
from keras.optimizer_v2 import gradient_descent
from keras.utils import dataset_creator
from tensorflow.python.platform import tf_logging as logging
class DatasetCreatorModelFitTestBase(tf.test.TestCase, parameterized.TestCase):
"""The base class for DatasetCreator with Model.fit tests."""
def _get_dataset_fn(self, use_lookup_layer):
if use_lookup_layer:
filepath = os.path.join(self.get_temp_dir(), "vocab")
with open(filepath, "w") as f:
f.write("\n".join(["earth", "wind", "and", "fire"]))
def dataset_fn(input_context):
del input_context
lookup_layer = string_lookup.StringLookup(
num_oov_indices=1, vocabulary=filepath)
x = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
y = np.array([0, 1])
map_fn = lambda x, y: (lookup_layer(x), y)
return tf.data.Dataset.from_tensor_slices(
(x, y)).shuffle(10).repeat().batch(2).map(map_fn)
else:
def dataset_fn(input_context):
del input_context
x = tf.random.uniform((10, 10))
y = tf.random.uniform((10,))
return tf.data.Dataset.from_tensor_slices(
(x, y)).shuffle(10).repeat().batch(2)
return dataset_fn
def _model_compile(self,
strategy,
steps_per_execution=1,
run_eagerly=False,
with_normalization_layer=False,
use_lookup_layer=False):
class ResultAssertingCallback(callbacks_lib.Callback):
"""A callback that asserts the result of the tests."""
def __init__(self):
self._prev_epoch = -1
def on_epoch_end(self, epoch, logs=None):
logging.info("testModelFit: epoch=%r, logs=%r", epoch, logs)
if epoch <= self._prev_epoch:
raise RuntimeError("Epoch is supposed to be larger than previous.")
self._prev_epoch = epoch
is_loss_float = (
logs.get("loss", None) is not None and
isinstance(logs["loss"], (float, np.floating)))
if not is_loss_float:
raise RuntimeError("loss is supposed to be in the logs and float.")
with strategy.scope():
model = sequential.Sequential([core_layers.Dense(10)])
if with_normalization_layer:
norm = keras.layers.BatchNormalization(
axis=-1, input_shape=(4, 4, 3), momentum=0.8)
model.add(norm)
model.add(core_layers.Dense(1, activation="sigmoid"))
self._accuracy_metric = keras.metrics.Accuracy()
model.compile(
gradient_descent.SGD(),
loss="binary_crossentropy",
metrics=[self._accuracy_metric],
steps_per_execution=steps_per_execution,
run_eagerly=run_eagerly)
return model, [ResultAssertingCallback()]
def _model_fit(self,
strategy,
steps_per_execution=1,
validation_data=None,
x=None,
y=None,
shuffle=True,
batch_size=None,
steps_per_epoch=10,
run_eagerly=False,
with_normalization_layer=False,
callbacks=None,
use_lookup_layer=False):
if callbacks is None:
callbacks = []
model, default_callbacks = self._model_compile(strategy,
steps_per_execution,
run_eagerly,
with_normalization_layer,
use_lookup_layer)
callbacks += default_callbacks
if x is None:
x = dataset_creator.DatasetCreator(self._get_dataset_fn(use_lookup_layer))
if validation_data is None:
validation_data = dataset_creator.DatasetCreator(
self._get_dataset_fn(use_lookup_layer))
model.fit(
x,
y,
shuffle=shuffle,
batch_size=batch_size,
epochs=10,
steps_per_epoch=steps_per_epoch,
callbacks=callbacks,
validation_data=validation_data,
validation_steps=steps_per_epoch)
return model
def _model_evaluate(self,
strategy,
steps_per_execution=1,
x=None,
y=None,
batch_size=None,
steps=10,
run_eagerly=False,
with_normalization_layer=False,
callbacks=None):
if callbacks is None:
callbacks = []
model, default_callbacks = self._model_compile(
strategy,
steps_per_execution,
run_eagerly,
with_normalization_layer,
)
callbacks += default_callbacks
def dataset_fn(input_context):
del input_context
x = tf.random.uniform((10, 10))
y = tf.random.uniform((10, 1))
return tf.data.Dataset.from_tensor_slices(
(x, y)).shuffle(10).repeat().batch(8)
if x is None:
x = dataset_creator.DatasetCreator(dataset_fn)
model.evaluate(
x=x, y=y, steps=steps, callbacks=callbacks, batch_size=batch_size)
return model
def _model_predict(
self,
strategy,
model=None,
steps_per_execution=1,
test_data=None,
steps=10,
with_normalization_layer=False,
):
callbacks = []
if model is None:
model, default_callbacks = self._model_compile(
strategy,
steps_per_execution,
with_normalization_layer=with_normalization_layer,
)
callbacks += default_callbacks
def create_test_data():
x = tf.constant([1., 2., 3., 1., 5., 1.])
return tf.data.Dataset.from_tensor_slices(x).repeat().batch(2)
if test_data is None:
test_data = create_test_data()
predictions = model.predict(x=test_data, steps=steps, callbacks=callbacks)
predictions = np.around(predictions, 4)
return model, predictions
Classes
class DatasetCreatorModelFitTestBase (methodName='runTest')
-
The base class for DatasetCreator with Model.fit tests.
Create an instance of the class that will use the named test method when executed. Raises a ValueError if the instance does not have a method with the specified name.
Expand source code
class DatasetCreatorModelFitTestBase(tf.test.TestCase, parameterized.TestCase): """The base class for DatasetCreator with Model.fit tests.""" def _get_dataset_fn(self, use_lookup_layer): if use_lookup_layer: filepath = os.path.join(self.get_temp_dir(), "vocab") with open(filepath, "w") as f: f.write("\n".join(["earth", "wind", "and", "fire"])) def dataset_fn(input_context): del input_context lookup_layer = string_lookup.StringLookup( num_oov_indices=1, vocabulary=filepath) x = np.array([["earth", "wind", "and", "fire"], ["fire", "and", "earth", "michigan"]]) y = np.array([0, 1]) map_fn = lambda x, y: (lookup_layer(x), y) return tf.data.Dataset.from_tensor_slices( (x, y)).shuffle(10).repeat().batch(2).map(map_fn) else: def dataset_fn(input_context): del input_context x = tf.random.uniform((10, 10)) y = tf.random.uniform((10,)) return tf.data.Dataset.from_tensor_slices( (x, y)).shuffle(10).repeat().batch(2) return dataset_fn def _model_compile(self, strategy, steps_per_execution=1, run_eagerly=False, with_normalization_layer=False, use_lookup_layer=False): class ResultAssertingCallback(callbacks_lib.Callback): """A callback that asserts the result of the tests.""" def __init__(self): self._prev_epoch = -1 def on_epoch_end(self, epoch, logs=None): logging.info("testModelFit: epoch=%r, logs=%r", epoch, logs) if epoch <= self._prev_epoch: raise RuntimeError("Epoch is supposed to be larger than previous.") self._prev_epoch = epoch is_loss_float = ( logs.get("loss", None) is not None and isinstance(logs["loss"], (float, np.floating))) if not is_loss_float: raise RuntimeError("loss is supposed to be in the logs and float.") with strategy.scope(): model = sequential.Sequential([core_layers.Dense(10)]) if with_normalization_layer: norm = keras.layers.BatchNormalization( axis=-1, input_shape=(4, 4, 3), momentum=0.8) model.add(norm) model.add(core_layers.Dense(1, activation="sigmoid")) self._accuracy_metric = keras.metrics.Accuracy() model.compile( gradient_descent.SGD(), loss="binary_crossentropy", metrics=[self._accuracy_metric], steps_per_execution=steps_per_execution, run_eagerly=run_eagerly) return model, [ResultAssertingCallback()] def _model_fit(self, strategy, steps_per_execution=1, validation_data=None, x=None, y=None, shuffle=True, batch_size=None, steps_per_epoch=10, run_eagerly=False, with_normalization_layer=False, callbacks=None, use_lookup_layer=False): if callbacks is None: callbacks = [] model, default_callbacks = self._model_compile(strategy, steps_per_execution, run_eagerly, with_normalization_layer, use_lookup_layer) callbacks += default_callbacks if x is None: x = dataset_creator.DatasetCreator(self._get_dataset_fn(use_lookup_layer)) if validation_data is None: validation_data = dataset_creator.DatasetCreator( self._get_dataset_fn(use_lookup_layer)) model.fit( x, y, shuffle=shuffle, batch_size=batch_size, epochs=10, steps_per_epoch=steps_per_epoch, callbacks=callbacks, validation_data=validation_data, validation_steps=steps_per_epoch) return model def _model_evaluate(self, strategy, steps_per_execution=1, x=None, y=None, batch_size=None, steps=10, run_eagerly=False, with_normalization_layer=False, callbacks=None): if callbacks is None: callbacks = [] model, default_callbacks = self._model_compile( strategy, steps_per_execution, run_eagerly, with_normalization_layer, ) callbacks += default_callbacks def dataset_fn(input_context): del input_context x = tf.random.uniform((10, 10)) y = tf.random.uniform((10, 1)) return tf.data.Dataset.from_tensor_slices( (x, y)).shuffle(10).repeat().batch(8) if x is None: x = dataset_creator.DatasetCreator(dataset_fn) model.evaluate( x=x, y=y, steps=steps, callbacks=callbacks, batch_size=batch_size) return model def _model_predict( self, strategy, model=None, steps_per_execution=1, test_data=None, steps=10, with_normalization_layer=False, ): callbacks = [] if model is None: model, default_callbacks = self._model_compile( strategy, steps_per_execution, with_normalization_layer=with_normalization_layer, ) callbacks += default_callbacks def create_test_data(): x = tf.constant([1., 2., 3., 1., 5., 1.]) return tf.data.Dataset.from_tensor_slices(x).repeat().batch(2) if test_data is None: test_data = create_test_data() predictions = model.predict(x=test_data, steps=steps, callbacks=callbacks) predictions = np.around(predictions, 4) return model, predictions
Ancestors
- tensorflow.python.framework.test_util.TensorFlowTestCase
- absl.testing.parameterized.TestCase
- absl.testing.absltest.TestCase
- absl.third_party.unittest3_backport.case.TestCase
- unittest.case.TestCase