268 lines
8.3 KiB
Python
268 lines
8.3 KiB
Python
# 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 os
|
|
|
|
import numpy as np
|
|
import tensorflow.compat.v2 as tf
|
|
from absl.testing import parameterized
|
|
|
|
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.optimizers.legacy import gradient_descent
|
|
from keras.utils import dataset_creator
|
|
|
|
# isort: off
|
|
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,
|
|
jit_compile=None,
|
|
):
|
|
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,
|
|
jit_compile=jit_compile,
|
|
)
|
|
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,
|
|
use_dataset_creator=True,
|
|
verbose="auto",
|
|
jit_compile=None,
|
|
):
|
|
if callbacks is None:
|
|
callbacks = []
|
|
|
|
model, default_callbacks = self._model_compile(
|
|
strategy,
|
|
steps_per_execution,
|
|
run_eagerly,
|
|
with_normalization_layer,
|
|
jit_compile,
|
|
)
|
|
callbacks += default_callbacks
|
|
|
|
if x is None:
|
|
if use_dataset_creator:
|
|
x = dataset_creator.DatasetCreator(
|
|
self._get_dataset_fn(use_lookup_layer)
|
|
)
|
|
else:
|
|
x = self._get_dataset_fn(use_lookup_layer)(None)
|
|
|
|
if validation_data is None:
|
|
if use_dataset_creator:
|
|
validation_data = dataset_creator.DatasetCreator(
|
|
self._get_dataset_fn(use_lookup_layer)
|
|
)
|
|
else:
|
|
validation_data = self._get_dataset_fn(use_lookup_layer)(None)
|
|
|
|
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,
|
|
verbose=verbose,
|
|
)
|
|
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,
|
|
use_dataset_creator=True,
|
|
):
|
|
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:
|
|
if use_dataset_creator:
|
|
x = dataset_creator.DatasetCreator(dataset_fn)
|
|
else:
|
|
x = dataset_fn(None)
|
|
|
|
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.0], [2.0], [3.0], [1.0], [5.0], [1.0]])
|
|
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
|