3069 lines
118 KiB
Python
3069 lines
118 KiB
Python
![]() |
# Copyright 2016 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 tf.keras models using tf.distribute.Strategy."""
|
||
|
|
||
|
import os
|
||
|
|
||
|
import numpy as np
|
||
|
import tensorflow.compat.v2 as tf
|
||
|
from absl.testing import parameterized
|
||
|
|
||
|
import keras
|
||
|
from keras import backend
|
||
|
from keras.distribute import distributed_training_utils
|
||
|
from keras.distribute import distributed_training_utils_v1
|
||
|
from keras.distribute import multi_worker_testing_utils
|
||
|
from keras.distribute import optimizer_combinations
|
||
|
from keras.distribute.strategy_combinations import all_strategies
|
||
|
from keras.distribute.strategy_combinations import (
|
||
|
multi_worker_mirrored_strategies,
|
||
|
)
|
||
|
from keras.distribute.strategy_combinations import (
|
||
|
strategies_minus_default_minus_tpu,
|
||
|
)
|
||
|
from keras.distribute.strategy_combinations import strategies_minus_tpu
|
||
|
from keras.distribute.strategy_combinations import tpu_strategies
|
||
|
from keras.engine import base_layer_utils
|
||
|
from keras.mixed_precision import policy
|
||
|
from keras.optimizers import optimizer as optimizer_base
|
||
|
from keras.optimizers.legacy import gradient_descent as gradient_descent_keras
|
||
|
from keras.testing_infra import test_utils
|
||
|
from keras.utils import losses_utils
|
||
|
from keras.utils import np_utils
|
||
|
|
||
|
# isort: off
|
||
|
from tensorflow.python.distribute.cluster_resolver import (
|
||
|
SimpleClusterResolver,
|
||
|
)
|
||
|
|
||
|
_RANDOM_SEED = 1337
|
||
|
_TRAIN_SIZE = 200
|
||
|
_INPUT_SIZE = (10,)
|
||
|
_NUM_CLASS = 2
|
||
|
|
||
|
# Note: Please make sure the tests in this file are also covered in
|
||
|
# keras_backward_compat_test for features that are supported with both APIs.
|
||
|
|
||
|
# TODO(anjalisridhar): Add a decorator that will allow us to run these tests as
|
||
|
# part of the tf.keras unit tests suite.
|
||
|
|
||
|
|
||
|
def simple_sequential_model():
|
||
|
model = keras.models.Sequential()
|
||
|
model.add(
|
||
|
keras.layers.Dense(16, activation="relu", input_shape=_INPUT_SIZE)
|
||
|
)
|
||
|
model.add(keras.layers.Dropout(0.1))
|
||
|
model.add(keras.layers.Dense(_NUM_CLASS, activation="softmax"))
|
||
|
return model
|
||
|
|
||
|
|
||
|
def simple_subclassed_model(num_labels=_NUM_CLASS):
|
||
|
class _SimpleMLP(keras.Model):
|
||
|
def __init__(self, num_labels):
|
||
|
super().__init__()
|
||
|
self.dense = keras.layers.Dense(num_labels)
|
||
|
|
||
|
def call(self, inputs):
|
||
|
return self.dense(inputs)
|
||
|
|
||
|
return _SimpleMLP(num_labels)
|
||
|
|
||
|
|
||
|
def simple_multi_inputs_multi_outputs_model():
|
||
|
input_a = keras.layers.Input(shape=(16,), name="input_a")
|
||
|
input_b = keras.layers.Input(shape=(16,), name="input_b")
|
||
|
|
||
|
merged = keras.layers.concatenate([input_a, input_b], name="merge")
|
||
|
output_c = keras.layers.Dense(3, activation="softmax", name="dense_2")(
|
||
|
merged
|
||
|
)
|
||
|
output_d = keras.layers.Dense(2, activation="softmax", name="dense_3")(
|
||
|
merged
|
||
|
)
|
||
|
model = keras.models.Model(
|
||
|
inputs=[input_a, input_b], outputs=[output_c, output_d]
|
||
|
)
|
||
|
return model
|
||
|
|
||
|
|
||
|
def get_multi_inputs_multi_outputs_data():
|
||
|
(a_train, c_train), (a_test, c_test) = test_utils.get_test_data(
|
||
|
train_samples=_TRAIN_SIZE,
|
||
|
test_samples=50,
|
||
|
input_shape=(16,),
|
||
|
num_classes=3,
|
||
|
random_seed=_RANDOM_SEED,
|
||
|
)
|
||
|
(b_train, d_train), (b_test, d_test) = test_utils.get_test_data(
|
||
|
train_samples=_TRAIN_SIZE,
|
||
|
test_samples=50,
|
||
|
input_shape=(16,),
|
||
|
num_classes=2,
|
||
|
random_seed=_RANDOM_SEED,
|
||
|
)
|
||
|
(m_train, _), (m_test, _) = test_utils.get_test_data(
|
||
|
train_samples=_TRAIN_SIZE,
|
||
|
test_samples=50,
|
||
|
input_shape=(8,),
|
||
|
num_classes=2,
|
||
|
random_seed=_RANDOM_SEED,
|
||
|
)
|
||
|
|
||
|
c_train = np_utils.to_categorical(c_train)
|
||
|
c_test = np_utils.to_categorical(c_test)
|
||
|
d_train = np_utils.to_categorical(d_train)
|
||
|
d_test = np_utils.to_categorical(d_test)
|
||
|
|
||
|
train_data = {
|
||
|
"input_a": a_train,
|
||
|
"input_b": b_train,
|
||
|
"input_m": m_train,
|
||
|
"output_c": c_train,
|
||
|
"output_d": d_train,
|
||
|
}
|
||
|
test_data = {
|
||
|
"input_a": a_test,
|
||
|
"input_b": b_test,
|
||
|
"input_m": m_test,
|
||
|
"output_c": c_test,
|
||
|
"output_d": d_test,
|
||
|
}
|
||
|
|
||
|
return (train_data, test_data)
|
||
|
|
||
|
|
||
|
def batch_wrapper(dataset, batch_size, distribution, repeat=None):
|
||
|
if repeat:
|
||
|
dataset = dataset.repeat(repeat)
|
||
|
# TPUs currently require fully defined input shapes, drop_remainder ensures
|
||
|
# the input will have fully defined shapes.
|
||
|
if backend.is_tpu_strategy(distribution):
|
||
|
return dataset.batch(batch_size, drop_remainder=True)
|
||
|
else:
|
||
|
return dataset.batch(batch_size)
|
||
|
|
||
|
|
||
|
def get_model():
|
||
|
x = keras.layers.Input(shape=(3,), name="input")
|
||
|
y = keras.layers.Dense(4, name="dense")(x)
|
||
|
model = keras.Model(x, y)
|
||
|
return model
|
||
|
|
||
|
|
||
|
def get_sample_weights_model():
|
||
|
x = keras.layers.Input(shape=(1,), name="input")
|
||
|
y = keras.layers.Dense(
|
||
|
1, kernel_initializer="ones", bias_initializer="zeros", name="dense"
|
||
|
)(x)
|
||
|
model = keras.Model(x, y)
|
||
|
return model
|
||
|
|
||
|
|
||
|
def get_dataset(distribution):
|
||
|
inputs = np.zeros((10, 3), dtype=np.float32)
|
||
|
targets = np.zeros((10, 4), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
|
||
|
dataset = dataset.repeat(100)
|
||
|
dataset = batch_wrapper(dataset, 10, distribution)
|
||
|
return dataset
|
||
|
|
||
|
|
||
|
def get_predict_dataset(distribution):
|
||
|
inputs = np.zeros((10, 3), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices(inputs)
|
||
|
dataset = dataset.repeat(100)
|
||
|
dataset = batch_wrapper(dataset, 10, distribution)
|
||
|
return dataset
|
||
|
|
||
|
|
||
|
def convert_numpy_to_dataset_with_unknown_cardinality(inputs, targets=None):
|
||
|
if targets is not None:
|
||
|
input_slices = (inputs, targets)
|
||
|
dummy_op = lambda inp, target: True
|
||
|
else:
|
||
|
input_slices = inputs
|
||
|
dummy_op = lambda inp: True
|
||
|
|
||
|
original_dataset = tf.data.Dataset.from_tensor_slices(input_slices)
|
||
|
ds_with_unknown_cardinality = original_dataset.filter(dummy_op).batch(
|
||
|
10, drop_remainder=True
|
||
|
)
|
||
|
return ds_with_unknown_cardinality
|
||
|
|
||
|
|
||
|
def multi_input_output_model():
|
||
|
a = keras.layers.Input(shape=(3,), name="input_a")
|
||
|
b = keras.layers.Input(shape=(5,), name="input_b")
|
||
|
# TODO(anjalisridhar): Change the output dimension of the second Dense layer
|
||
|
# once the iterator output validation issue has been fixed.
|
||
|
dense_1 = keras.layers.Dense(7, name="dense_1")
|
||
|
dense_2 = keras.layers.Dense(7, name="dense_2")
|
||
|
c = dense_1(a)
|
||
|
d = dense_2(b)
|
||
|
e = keras.layers.Dropout(0.5, name="dropout")(c)
|
||
|
model = keras.models.Model([a, b], [d, e])
|
||
|
return model
|
||
|
|
||
|
|
||
|
def strategy_minus_tpu_combinations():
|
||
|
return tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_tpu, mode=["graph", "eager"]
|
||
|
)
|
||
|
|
||
|
|
||
|
def tpu_strategy_combinations():
|
||
|
return tf.__internal__.test.combinations.combine(
|
||
|
distribution=tpu_strategies, mode=["graph", "eager"]
|
||
|
)
|
||
|
|
||
|
|
||
|
def tpu_strategy_combinations_graph_only():
|
||
|
return tf.__internal__.test.combinations.combine(
|
||
|
distribution=tpu_strategies, mode=["graph"]
|
||
|
)
|
||
|
|
||
|
|
||
|
def multi_worker_strategy_combinations_eager_only():
|
||
|
return tf.__internal__.test.combinations.combine(
|
||
|
distribution=multi_worker_mirrored_strategies, mode=["eager"]
|
||
|
)
|
||
|
|
||
|
|
||
|
def all_strategy_combinations():
|
||
|
return (
|
||
|
strategy_minus_tpu_combinations()
|
||
|
+ tpu_strategy_combinations()
|
||
|
+ multi_worker_strategy_combinations_eager_only()
|
||
|
)
|
||
|
|
||
|
|
||
|
def all_strategy_minus_default_and_tpu_combinations():
|
||
|
return tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy_gpu,
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus, # noqa: E501
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
)
|
||
|
|
||
|
|
||
|
def all_strategy_combinations_minus_default():
|
||
|
return (
|
||
|
all_strategy_minus_default_and_tpu_combinations()
|
||
|
+ tpu_strategy_combinations()
|
||
|
+ multi_worker_strategy_combinations_eager_only()
|
||
|
)
|
||
|
|
||
|
|
||
|
def strategy_and_optimizer_combinations():
|
||
|
non_tpu_strategies = tf.__internal__.test.combinations.times(
|
||
|
strategy_minus_tpu_combinations(),
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
optimizer=[
|
||
|
optimizer_combinations.adagrad_optimizer_v1_fn,
|
||
|
optimizer_combinations.adam_optimizer_v1_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_v1_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_v1_fn,
|
||
|
optimizer_combinations.adadelta_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adagrad_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adamax_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.nadam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.ftrl_optimizer_keras_v2_fn,
|
||
|
]
|
||
|
),
|
||
|
)
|
||
|
tpu_strategies_graph = tf.__internal__.test.combinations.combine(
|
||
|
distribution=tpu_strategies,
|
||
|
mode=["graph"],
|
||
|
optimizer=[
|
||
|
optimizer_combinations.adagrad_optimizer_v1_fn,
|
||
|
optimizer_combinations.adam_optimizer_v1_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_v1_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_v1_fn,
|
||
|
optimizer_combinations.adagrad_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
|
||
|
],
|
||
|
)
|
||
|
tpu_strategies_eager = tf.__internal__.test.combinations.combine(
|
||
|
distribution=tpu_strategies,
|
||
|
mode=["eager"],
|
||
|
optimizer=[
|
||
|
optimizer_combinations.adagrad_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
|
||
|
],
|
||
|
)
|
||
|
multi_worker_eager = tf.__internal__.test.combinations.combine(
|
||
|
distribution=multi_worker_mirrored_strategies,
|
||
|
mode=["eager"],
|
||
|
optimizer=[
|
||
|
optimizer_combinations.adadelta_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adagrad_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.adamax_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.nadam_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
|
||
|
optimizer_combinations.ftrl_optimizer_keras_v2_fn,
|
||
|
],
|
||
|
)
|
||
|
return (
|
||
|
non_tpu_strategies
|
||
|
+ tpu_strategies_eager
|
||
|
+ tpu_strategies_graph
|
||
|
+ multi_worker_eager
|
||
|
)
|
||
|
|
||
|
|
||
|
class BatchCountingCB(keras.callbacks.Callback):
|
||
|
def __init__(self):
|
||
|
super().__init__()
|
||
|
self.train_begin_batches = []
|
||
|
self.train_end_batches = []
|
||
|
self.test_begin_batches = []
|
||
|
self.test_end_batches = []
|
||
|
self.predict_begin_batches = []
|
||
|
self.predict_end_batches = []
|
||
|
|
||
|
def on_train_batch_begin(self, batch, logs=None):
|
||
|
self.train_begin_batches.append(batch)
|
||
|
|
||
|
def on_train_batch_end(self, batch, logs=None):
|
||
|
self.train_end_batches.append(batch)
|
||
|
|
||
|
def on_test_batch_begin(self, batch, logs=None):
|
||
|
self.test_begin_batches.append(batch)
|
||
|
|
||
|
def on_test_batch_end(self, batch, logs=None):
|
||
|
self.test_end_batches.append(batch)
|
||
|
|
||
|
def on_predict_batch_begin(self, batch, logs=None):
|
||
|
self.predict_begin_batches.append(batch)
|
||
|
|
||
|
def on_predict_batch_end(self, batch, logs=None):
|
||
|
self.predict_end_batches.append(batch)
|
||
|
|
||
|
|
||
|
class TestDistributionStrategyWithNumpyArrays(
|
||
|
tf.test.TestCase, parameterized.TestCase
|
||
|
):
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calculating_input_params_no_steps_no_batch_size(
|
||
|
self, distribution
|
||
|
):
|
||
|
# Calculate the per_replica_batch_size scaling factor for strategies
|
||
|
# that use per_core_batch_size
|
||
|
replica_scale_factor = 1.0
|
||
|
if not distributed_training_utils.global_batch_size_supported(
|
||
|
distribution
|
||
|
):
|
||
|
replica_scale_factor = distribution.num_replicas_in_sync
|
||
|
|
||
|
with self.cached_session():
|
||
|
# Default global batch size 32 for input with 64 samples run in 2
|
||
|
# steps
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=None, batch_size=None
|
||
|
)
|
||
|
self.assertEqual(batch_size, 32 // replica_scale_factor)
|
||
|
self.assertEqual(steps, 2)
|
||
|
|
||
|
# Computed global batch size 20 is lower than 32 if we pass less
|
||
|
# samples.
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 20, steps=None, batch_size=None
|
||
|
)
|
||
|
self.assertEqual(batch_size, 20 // replica_scale_factor)
|
||
|
self.assertEqual(steps, 1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calculating_input_params_with_steps_no_batch_size(
|
||
|
self, distribution
|
||
|
):
|
||
|
# Calculate the per_replica_batch_size scaling factor for strategies
|
||
|
# that use per_core_batch_size
|
||
|
replica_scale_factor = 1.0
|
||
|
if not distributed_training_utils.global_batch_size_supported(
|
||
|
distribution
|
||
|
):
|
||
|
replica_scale_factor = distribution.num_replicas_in_sync
|
||
|
|
||
|
with self.cached_session():
|
||
|
# Computed global batch size is correct for number of specified 1
|
||
|
# step
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=1, batch_size=None
|
||
|
)
|
||
|
self.assertEqual(batch_size, 64 // replica_scale_factor)
|
||
|
self.assertEqual(steps, 1)
|
||
|
|
||
|
# Computed global batch size is correct for number of specified 2
|
||
|
# steps
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=2, batch_size=None
|
||
|
)
|
||
|
self.assertEqual(batch_size, 32 // replica_scale_factor)
|
||
|
self.assertEqual(steps, 2)
|
||
|
|
||
|
# All samples can not be consumed in specified number of steps
|
||
|
with self.assertRaisesRegex(ValueError, "not divisible by steps"):
|
||
|
distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 63, steps=2, batch_size=None
|
||
|
)
|
||
|
|
||
|
# This cases is different for different strategies due to the
|
||
|
# difference in supported batch size being global or per-replica.
|
||
|
if replica_scale_factor == 1:
|
||
|
# Computed global batch size is correct even if not sharadable
|
||
|
(
|
||
|
steps,
|
||
|
batch_size,
|
||
|
) = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 63, steps=3, batch_size=None
|
||
|
)
|
||
|
self.assertEqual(batch_size, 21)
|
||
|
self.assertEqual(steps, 3)
|
||
|
else:
|
||
|
# Computed global batch size can not be sharded across replicas
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError,
|
||
|
"could not be sharded evenly across the sync replicas",
|
||
|
):
|
||
|
distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 63, steps=1, batch_size=None
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calculating_input_params_no_steps_with_batch_size(
|
||
|
self, distribution
|
||
|
):
|
||
|
# Calculate the per_replica_batch_size scaling factor for strategies
|
||
|
# that use per_core_batch_size
|
||
|
replica_scale_factor = 1.0
|
||
|
if not distributed_training_utils.global_batch_size_supported(
|
||
|
distribution
|
||
|
):
|
||
|
replica_scale_factor = distribution.num_replicas_in_sync
|
||
|
|
||
|
with self.cached_session():
|
||
|
# Computed steps is correct for specified batch size
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=None, batch_size=16
|
||
|
)
|
||
|
self.assertEqual(batch_size, 16)
|
||
|
self.assertEqual(steps, 4 // replica_scale_factor)
|
||
|
|
||
|
# Computed steps is correct for specified batch size
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=None, batch_size=32
|
||
|
)
|
||
|
self.assertEqual(batch_size, 32)
|
||
|
self.assertEqual(steps, 2 // replica_scale_factor)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calculating_input_params_with_steps_with_batch_size(
|
||
|
self, distribution
|
||
|
):
|
||
|
with self.cached_session():
|
||
|
# No change to steps and batch size if both specified and feasible
|
||
|
steps, batch_size = distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=5, batch_size=3
|
||
|
)
|
||
|
self.assertEqual(batch_size, 3)
|
||
|
self.assertEqual(steps, 5)
|
||
|
|
||
|
# Number of samples is less than global batch size * steps
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError, "less than samples required"
|
||
|
):
|
||
|
distributed_training_utils_v1.get_input_params(
|
||
|
distribution, 64, steps=10, batch_size=13
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calling_model_with_numpy_arrays(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae"]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
inputs = np.zeros((64, 3), dtype=np.float32)
|
||
|
targets = np.zeros((64, 4), dtype=np.float32)
|
||
|
|
||
|
# Call fit with validation data
|
||
|
model.fit(
|
||
|
inputs,
|
||
|
targets,
|
||
|
epochs=1,
|
||
|
batch_size=2,
|
||
|
verbose=0,
|
||
|
validation_data=(inputs, targets),
|
||
|
)
|
||
|
|
||
|
# TODO(anjalisridhar): We need tests for when the batch size and
|
||
|
# steps are smaller and results in a 0 batch_size and steps
|
||
|
# value.
|
||
|
model.evaluate(inputs, targets)
|
||
|
model.evaluate(inputs, targets, batch_size=8)
|
||
|
|
||
|
model.predict(inputs)
|
||
|
model.predict(inputs, batch_size=8)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calling_model_with_mixed_precision(self, distribution):
|
||
|
if isinstance(
|
||
|
distribution,
|
||
|
(
|
||
|
tf.compat.v1.distribute.experimental.ParameterServerStrategy,
|
||
|
tf.distribute.experimental.ParameterServerStrategy,
|
||
|
tf.distribute.experimental.CentralStorageStrategy,
|
||
|
tf.compat.v1.distribute.experimental.CentralStorageStrategy,
|
||
|
),
|
||
|
):
|
||
|
self.skipTest("b/152097775")
|
||
|
if backend.is_tpu_strategy(distribution):
|
||
|
policy_name = "mixed_bfloat16"
|
||
|
else:
|
||
|
policy_name = "mixed_float16"
|
||
|
with self.cached_session(), distribution.scope(), policy.policy_scope(
|
||
|
policy_name
|
||
|
):
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
x = keras.layers.Input(shape=(3,), name="input")
|
||
|
y = keras.layers.Dense(4, name="dense")(x)
|
||
|
y = keras.layers.Activation("softmax", dtype="float32")(y)
|
||
|
model = keras.Model(x, y)
|
||
|
loss = "mse"
|
||
|
metrics = ["mae"]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
# We need to pass float32 since TPUs do not support float64, even
|
||
|
# though these arrays will immediately be casted to bfloat16 on
|
||
|
# TPUs. We also cannot pass bfloat16, as Numpy does not support it.
|
||
|
inputs = np.zeros((64, 3), dtype="float32")
|
||
|
targets = np.zeros((64, 4), dtype="float32")
|
||
|
|
||
|
model.fit(
|
||
|
inputs,
|
||
|
targets,
|
||
|
epochs=1,
|
||
|
batch_size=2,
|
||
|
verbose=0,
|
||
|
validation_data=(inputs, targets),
|
||
|
)
|
||
|
|
||
|
model.evaluate(inputs, targets)
|
||
|
model.evaluate(inputs, targets, batch_size=8)
|
||
|
|
||
|
model.predict(inputs)
|
||
|
model.predict(inputs, batch_size=8)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_operator_overload_mixed_precision(self, distribution):
|
||
|
# Regression test that tests a fixed bug does not reoccur. Adding an
|
||
|
# AutoCastVariable to a tensor on a TPU, where the variable was the LHS
|
||
|
# of the '+' operator, used to cause the gradient w.r.t. the variable to
|
||
|
# be None.
|
||
|
if isinstance(
|
||
|
distribution,
|
||
|
(
|
||
|
tf.compat.v1.distribute.experimental.ParameterServerStrategy,
|
||
|
tf.distribute.experimental.ParameterServerStrategy,
|
||
|
tf.distribute.experimental.CentralStorageStrategy,
|
||
|
tf.compat.v1.distribute.experimental.CentralStorageStrategy,
|
||
|
),
|
||
|
):
|
||
|
self.skipTest("b/152097775")
|
||
|
|
||
|
if backend.is_tpu_strategy(distribution):
|
||
|
policy_name = "mixed_bfloat16"
|
||
|
else:
|
||
|
policy_name = "mixed_float16"
|
||
|
|
||
|
class MyLayer(keras.layers.Layer):
|
||
|
def build(self, _):
|
||
|
self.v1 = self.add_weight("v", ())
|
||
|
self.v2 = self.add_weight("v", ())
|
||
|
|
||
|
def call(self, inp):
|
||
|
inp += self.v1
|
||
|
return self.v2 + inp
|
||
|
|
||
|
with self.cached_session(), distribution.scope():
|
||
|
layer = MyLayer(dtype=policy_name)
|
||
|
|
||
|
def run_fn():
|
||
|
x = np.array([1.0])
|
||
|
with tf.GradientTape() as tape:
|
||
|
y = layer(x)
|
||
|
grad_v1, grad_v2 = tape.gradient(y, [layer.v1, layer.v2])
|
||
|
return grad_v1, grad_v2
|
||
|
|
||
|
if tf.executing_eagerly():
|
||
|
run_fn = tf.function(run_fn)
|
||
|
|
||
|
grad_v1, grad_v2 = distribution.run(run_fn)
|
||
|
self.assertIsNotNone(grad_v1)
|
||
|
self.assertIsNotNone(grad_v2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_optimizer_in_cross_replica_context_raises_error(
|
||
|
self, distribution
|
||
|
):
|
||
|
|
||
|
with self.cached_session(), distribution.scope():
|
||
|
model = keras.models.Sequential([keras.layers.Dense(1)])
|
||
|
x = np.array([[1.0]])
|
||
|
with tf.GradientTape() as tape:
|
||
|
y = model(x)
|
||
|
gradients = tape.gradient(y, model.trainable_variables)
|
||
|
optimizer = gradient_descent_keras.SGD()
|
||
|
|
||
|
with self.assertRaisesRegex(
|
||
|
RuntimeError, "cannot be called in cross-replica context"
|
||
|
):
|
||
|
optimizer.apply_gradients(
|
||
|
zip(gradients, model.trainable_variables)
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calling_model_with_nested_numpy_arrays(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = multi_input_output_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
input_a_np = np.asarray(np.random.random((64, 3)), dtype=np.float32)
|
||
|
input_b_np = np.asarray(np.random.random((64, 5)), dtype=np.float32)
|
||
|
inputs = [input_a_np, input_b_np]
|
||
|
|
||
|
output_d_np = np.asarray(
|
||
|
np.random.random((64, 7)), dtype=np.float32
|
||
|
)
|
||
|
output_e_np = np.asarray(
|
||
|
np.random.random((64, 7)), dtype=np.float32
|
||
|
)
|
||
|
targets = [output_d_np, output_e_np]
|
||
|
|
||
|
# Call fit with validation data
|
||
|
model.fit(inputs, targets, epochs=1, batch_size=8, verbose=0)
|
||
|
|
||
|
# TODO(anjalisridhar): We need tests for when the batch size and
|
||
|
# steps are smaller and results in a 0 batch_size and steps value.
|
||
|
model.evaluate(inputs, targets)
|
||
|
model.evaluate(inputs, targets, batch_size=8)
|
||
|
|
||
|
model.predict(inputs)
|
||
|
model.predict(inputs, batch_size=8)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_tpu, mode=["graph", "eager"]
|
||
|
)
|
||
|
+ tf.__internal__.test.combinations.combine(
|
||
|
distribution=multi_worker_mirrored_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_numpy_with_sample_weights(self, distribution):
|
||
|
with self.cached_session(), distribution.scope():
|
||
|
model = get_sample_weights_model()
|
||
|
optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.array([[0], [1], [2], [3]], np.float32)
|
||
|
targets = np.array([[2], [4], [6], [8]], np.float32)
|
||
|
sample_weights = np.array([0.25, 0.5, 0.75, 1], np.float32)
|
||
|
|
||
|
result = model.evaluate(
|
||
|
inputs,
|
||
|
targets,
|
||
|
batch_size=2,
|
||
|
sample_weight=sample_weights,
|
||
|
verbose=1,
|
||
|
)
|
||
|
|
||
|
# The per sample loss is multiplied by the corresponding sample
|
||
|
# weight. The average of these weighted losses is the return value
|
||
|
# of the `evaluate` call. For example, in the test above the average
|
||
|
# weighted loss is calculated in the following manner:
|
||
|
|
||
|
# batch_1 = (((2-0)^2) * 0.25 + ((4-1)^2) * 0.5) / 2 = 5.5 / 2 =
|
||
|
# 2.75
|
||
|
# batch_2 = (((6-2)^2 * 0.75) + ((8-3)^2 * 1)) / 2 = 37 / 2 = 18.5
|
||
|
# final result = (batch_1 + batch_2) / 2 = 10.625.
|
||
|
# The first time we divide by number of input samples and the second
|
||
|
# time we divide by number of steps/batches that the loss is
|
||
|
# aggregated over.
|
||
|
self.assertAllClose(result, 10.625)
|
||
|
|
||
|
# We now test without passing sample_weights:
|
||
|
# batch_1 = ((2-0)^2) + ((4-1)^2) / 2 = 13 / 2 = 6.5
|
||
|
# batch_2 = ((6-2)^2) + ((8-3)^2) / 2 = 41 / 2 = 20.5
|
||
|
# final result = (batch_1 + batch_2) / 2 = 27 / 2 = 13.5
|
||
|
result = model.evaluate(inputs, targets, batch_size=2, verbose=1)
|
||
|
self.assertAllClose(result, 13.5)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_flatten_predict_outputs(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
model = multi_input_output_model()
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
# We take 6 input samples with each input having a dimension of 3 or
|
||
|
# 5.
|
||
|
input_a_np = np.asarray(np.random.random((6, 3)), dtype=np.float32)
|
||
|
input_b_np = np.asarray(np.random.random((6, 5)), dtype=np.float32)
|
||
|
inputs = [input_a_np, input_b_np]
|
||
|
|
||
|
outs = model.predict(inputs)
|
||
|
# `predict` a list that is equal in length to the number of model
|
||
|
# outputs. In this test our model has two outputs and each element
|
||
|
# of `outs` corresponds to all the samples of one of the model
|
||
|
# outputs.
|
||
|
self.assertLen(outs, 2)
|
||
|
# Each of the output samples have a dimension of 7. We should
|
||
|
# process all the available input samples(6).
|
||
|
self.assertAllEqual([6, 7], outs[0].shape)
|
||
|
self.assertAllEqual([6, 7], outs[1].shape)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only(),
|
||
|
tf.__internal__.test.combinations.combine(batch_size=[4, 6]),
|
||
|
)
|
||
|
)
|
||
|
def test_evaluate_with_partial_batch(self, distribution, batch_size):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = get_model()
|
||
|
model_with_ds_strategy.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
cpu_model = get_model()
|
||
|
cpu_model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
x = np.random.random((10, 3)).astype("float32")
|
||
|
y = np.random.random((10, 4)).astype("float32")
|
||
|
|
||
|
# As sample size is 10, we batch by 4 so that the last batch is a
|
||
|
# partial batch. Also `evaluate()` using numpy array as inputs
|
||
|
# without distribution strategy uses entire sample as a single
|
||
|
# batch. As so, we remove parameters `batch_size` and `steps`.
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
evaluate_ground_truth = cpu_model.evaluate(x, y)
|
||
|
|
||
|
# We don't compare the loss as loss is currently not computed as
|
||
|
# metric in Keras, the loss value is inaccurate for last partial
|
||
|
# batch due to more weights for the last batch samples.
|
||
|
steps = np.ceil(10.0 / batch_size)
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.evaluate(
|
||
|
x, y, batch_size=batch_size, steps=steps
|
||
|
)[1:],
|
||
|
evaluate_ground_truth[1:],
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
# Test that `steps` is inferred correctly when final partial batch
|
||
|
# exists.
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.evaluate(x, y, batch_size=batch_size)[
|
||
|
1:
|
||
|
],
|
||
|
evaluate_ground_truth[1:],
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
)
|
||
|
def test_predict_with_partial_batch(self, distribution):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = get_model()
|
||
|
model_with_ds_strategy.compile(optimizer, loss)
|
||
|
|
||
|
cpu_model = get_model()
|
||
|
cpu_model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.random.random((10, 3)).astype(np.float32)
|
||
|
|
||
|
# As sample size is 10, we batch by 4 so that the last batch is
|
||
|
# a partial batch. Also `predict()` using numpy array as inputs
|
||
|
# without distribution strategy uses entire sample as a single
|
||
|
# batch. As so, we remove parameters `batch_size` and `steps`.
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
predict_ground_truth = cpu_model.predict(inputs)
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.predict(inputs, batch_size=4, steps=3),
|
||
|
predict_ground_truth,
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
# Test that `steps` is inferred correctly when final partial batch
|
||
|
# exists.
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.predict(inputs, batch_size=4),
|
||
|
predict_ground_truth,
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
def test_no_target_model(self, distribution):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
|
||
|
class MyLayer(keras.layers.Layer):
|
||
|
def call(self, inputs, training=None):
|
||
|
self.add_loss(tf.reduce_sum(inputs), inputs=True)
|
||
|
return inputs
|
||
|
|
||
|
with distribution.scope():
|
||
|
model = keras.models.Sequential()
|
||
|
model.add(
|
||
|
keras.layers.Dense(
|
||
|
16, activation="relu", input_shape=_INPUT_SIZE
|
||
|
)
|
||
|
)
|
||
|
model.add(MyLayer())
|
||
|
model.add(keras.layers.Dense(_NUM_CLASS, activation="softmax"))
|
||
|
|
||
|
model.compile(optimizer)
|
||
|
inputs = np.zeros((20, 10), np.float32)
|
||
|
|
||
|
model.fit(inputs, epochs=1, steps_per_epoch=2)
|
||
|
model.predict(inputs, steps=1)
|
||
|
model.evaluate(inputs, steps=1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
)
|
||
|
def test_predict_multi_output_model_with_partial_batch(self, distribution):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = (
|
||
|
simple_multi_inputs_multi_outputs_model()
|
||
|
)
|
||
|
model_with_ds_strategy.compile(optimizer, loss)
|
||
|
|
||
|
cpu_model = simple_multi_inputs_multi_outputs_model()
|
||
|
cpu_model.compile(optimizer, loss)
|
||
|
|
||
|
input_data, _ = get_multi_inputs_multi_outputs_data()
|
||
|
input_dict = {
|
||
|
"input_a": input_data["input_a"],
|
||
|
"input_b": input_data["input_b"],
|
||
|
}
|
||
|
|
||
|
# As sample size is 200, we batch by 18 so that the last batch is
|
||
|
# a partial batch. Also `fit()` using numpy array as inputs without
|
||
|
# distribution strategy uses entire sample as a single batch. As so,
|
||
|
# we remove parameters `batch_size` and `steps`.
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.predict(
|
||
|
input_dict, batch_size=18, steps=12
|
||
|
),
|
||
|
cpu_model.predict(input_dict),
|
||
|
atol=1e-4,
|
||
|
rtol=1e-4,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_gradients_are_none(self, distribution):
|
||
|
|
||
|
if not tf.executing_eagerly():
|
||
|
self.skipTest("None gradients are not supported in graph mode")
|
||
|
|
||
|
class DenseWithExtraWeight(keras.layers.Dense):
|
||
|
def build(self, input_shape):
|
||
|
# Gradients w.r.t. extra_weights are None
|
||
|
self.extra_weight_1 = self.add_weight(
|
||
|
"extra_weight_1", shape=(), initializer="ones"
|
||
|
)
|
||
|
super().build(input_shape)
|
||
|
self.extra_weight_2 = self.add_weight(
|
||
|
"extra_weight_2", shape=(), initializer="ones"
|
||
|
)
|
||
|
|
||
|
with distribution.scope():
|
||
|
model = keras.Sequential(
|
||
|
[DenseWithExtraWeight(4, input_shape=(4,))]
|
||
|
)
|
||
|
model.compile("adam", "mse")
|
||
|
|
||
|
inputs = np.random.normal(size=(64, 4))
|
||
|
targets = np.random.normal(size=(64, 4))
|
||
|
old_kernel = model.get_weights()[1]
|
||
|
model.fit(inputs, targets)
|
||
|
new_kernel = model.get_weights()[1]
|
||
|
self.assertNotAllEqual(old_kernel, new_kernel)
|
||
|
|
||
|
|
||
|
class TestDistributionStrategyWithDatasets(
|
||
|
tf.test.TestCase, parameterized.TestCase
|
||
|
):
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_calling_model_on_same_dataset(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
dataset = get_dataset(distribution)
|
||
|
|
||
|
# Call fit with validation data
|
||
|
model.fit(
|
||
|
dataset,
|
||
|
epochs=1,
|
||
|
steps_per_epoch=2,
|
||
|
verbose=0,
|
||
|
validation_data=dataset,
|
||
|
validation_steps=2,
|
||
|
)
|
||
|
model.fit(
|
||
|
dataset,
|
||
|
epochs=1,
|
||
|
steps_per_epoch=2,
|
||
|
verbose=0,
|
||
|
validation_data=dataset,
|
||
|
validation_steps=2,
|
||
|
)
|
||
|
model.predict(get_predict_dataset(distribution), steps=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_model_interleaved_eval_same_as_direct_eval(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
user_controlled_model = get_model()
|
||
|
user_controlled_model.compile(
|
||
|
optimizer_fn(0.001),
|
||
|
loss="mse",
|
||
|
metrics=["mae", keras.metrics.CategoricalAccuracy()],
|
||
|
)
|
||
|
|
||
|
interleaved_model = get_model()
|
||
|
interleaved_model.set_weights(
|
||
|
user_controlled_model.get_weights()
|
||
|
)
|
||
|
interleaved_model.compile(
|
||
|
optimizer_fn(0.001),
|
||
|
loss="mse",
|
||
|
metrics=["mae", keras.metrics.CategoricalAccuracy()],
|
||
|
)
|
||
|
|
||
|
dataset = get_dataset(distribution)
|
||
|
|
||
|
# Call fit with validation interleaved
|
||
|
interleaved_output = interleaved_model.fit(
|
||
|
dataset,
|
||
|
epochs=2,
|
||
|
steps_per_epoch=2,
|
||
|
verbose=1,
|
||
|
validation_data=dataset,
|
||
|
validation_steps=2,
|
||
|
shuffle=False,
|
||
|
)
|
||
|
|
||
|
# Manually control the validation running after each epoch.
|
||
|
user_controlled_output = []
|
||
|
for _ in range(2):
|
||
|
user_controlled_model.fit(
|
||
|
dataset,
|
||
|
epochs=1,
|
||
|
steps_per_epoch=2,
|
||
|
verbose=1,
|
||
|
shuffle=False,
|
||
|
)
|
||
|
user_controlled_output.append(
|
||
|
user_controlled_model.evaluate(dataset, steps=2)
|
||
|
)
|
||
|
|
||
|
self.assertEqual(
|
||
|
interleaved_output.history["val_loss"],
|
||
|
[x[0] for x in user_controlled_output],
|
||
|
)
|
||
|
val_mean_absolute_error = interleaved_output.history.get(
|
||
|
"val_mean_absolute_error"
|
||
|
)
|
||
|
if not val_mean_absolute_error:
|
||
|
# The name of the metric changed in TF2.0
|
||
|
val_mean_absolute_error = interleaved_output.history["val_mae"]
|
||
|
self.assertEqual(
|
||
|
val_mean_absolute_error, [x[1] for x in user_controlled_output]
|
||
|
)
|
||
|
self.assertEqual(
|
||
|
interleaved_output.history["val_categorical_accuracy"],
|
||
|
[x[2] for x in user_controlled_output],
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_fit_with_tuple_and_dict_dataset_inputs(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = multi_input_output_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
input_a_np = np.random.random((10, 3)).astype("float32")
|
||
|
input_b_np = np.random.random((10, 5)).astype("float32")
|
||
|
output_d_np = np.random.random((10, 7)).astype("float32")
|
||
|
output_e_np = np.random.random((10, 7)).astype("float32")
|
||
|
|
||
|
# Test with tuples
|
||
|
dataset_tuple = tf.data.Dataset.from_tensor_slices(
|
||
|
((input_a_np, input_b_np), (output_d_np, output_e_np))
|
||
|
)
|
||
|
dataset_tuple = dataset_tuple.repeat(100)
|
||
|
dataset_tuple = dataset_tuple.batch(10)
|
||
|
|
||
|
model.fit(dataset_tuple, epochs=1, steps_per_epoch=2, verbose=1)
|
||
|
|
||
|
# Test with dict
|
||
|
dataset_dict = tf.data.Dataset.from_tensor_slices(
|
||
|
(
|
||
|
{"input_a": input_a_np, "input_b": input_b_np},
|
||
|
(output_d_np, output_e_np),
|
||
|
)
|
||
|
)
|
||
|
dataset_dict = dataset_dict.repeat(100)
|
||
|
dataset_dict = dataset_dict.batch(10)
|
||
|
|
||
|
model.fit(dataset_dict, epochs=1, steps_per_epoch=2, verbose=1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_fit_with_dictionary_in_the_dataset_b135161171(self, distribution):
|
||
|
|
||
|
if backend.is_tpu_strategy(distribution):
|
||
|
self.skipTest("b/142805125")
|
||
|
|
||
|
def custom_loss(predict, label, weight):
|
||
|
bce = keras.losses.binary_crossentropy(label, predict)
|
||
|
return tf.reduce_mean(bce * weight)
|
||
|
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
input_img = keras.layers.Input([64, 64, 3], name="img")
|
||
|
input_lbl = keras.layers.Input([64, 64, 1], name="lbl")
|
||
|
input_weight = keras.layers.Input([64, 64], name="weight")
|
||
|
predict = keras.layers.Conv2D(2, [1, 1], padding="same")(
|
||
|
input_img
|
||
|
)
|
||
|
loss_lambda = keras.layers.Lambda(
|
||
|
lambda x: custom_loss(*x), name="my_loss"
|
||
|
)
|
||
|
my_loss = loss_lambda([predict, input_lbl, input_weight])
|
||
|
model = keras.models.Model(
|
||
|
inputs=[input_img, input_lbl, input_weight],
|
||
|
outputs=[predict, my_loss],
|
||
|
)
|
||
|
model.add_loss(model.get_layer("my_loss").output)
|
||
|
model.compile(optimizer="adam")
|
||
|
|
||
|
if tf.executing_eagerly():
|
||
|
|
||
|
def map_fn(img, lbl, weight):
|
||
|
inputs = {"img": img, "lbl": lbl, "weight": weight}
|
||
|
return (inputs,)
|
||
|
|
||
|
else:
|
||
|
|
||
|
def map_fn(img, lbl, weight):
|
||
|
inputs = {"img": img, "lbl": lbl, "weight": weight}
|
||
|
return inputs, {}
|
||
|
|
||
|
fake_imgs = np.ones([50, 64, 64, 3], dtype=np.float32)
|
||
|
fake_lbls = np.ones([50, 64, 64, 1], dtype=np.float32)
|
||
|
fake_weights = np.ones([50, 64, 64], dtype=np.float32)
|
||
|
|
||
|
data = (
|
||
|
tf.data.Dataset.from_tensor_slices(
|
||
|
(fake_imgs, fake_lbls, fake_weights)
|
||
|
)
|
||
|
.map(map_fn)
|
||
|
.batch(10)
|
||
|
)
|
||
|
|
||
|
model.fit(data)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_fit_eval_and_predict_methods_on_dataset_without_steps(
|
||
|
self, distribution
|
||
|
):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
inputs = np.zeros((1000, 3), dtype=np.float32)
|
||
|
targets = np.zeros((1000, 4), dtype=np.float32)
|
||
|
# steps/steps_per_epoch are calculated when using numpy arrays as
|
||
|
# input data.
|
||
|
fit_with_numpy = model.fit(
|
||
|
inputs, targets, epochs=1, batch_size=10
|
||
|
).history
|
||
|
eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
|
||
|
predict_with_numpy = model.predict(inputs, batch_size=10)
|
||
|
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
|
||
|
dataset = dataset.batch(10, drop_remainder=True)
|
||
|
fit_with_ds = model.fit(dataset, epochs=1).history
|
||
|
eval_with_ds = model.evaluate(dataset)
|
||
|
predict_dataset = tf.data.Dataset.from_tensor_slices(inputs)
|
||
|
predict_dataset = predict_dataset.batch(10, drop_remainder=True)
|
||
|
predict_with_ds = model.predict(predict_dataset)
|
||
|
self.assertAllClose(
|
||
|
fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_predict_on_dataset_with_unknown_cardinality_without_steps(
|
||
|
self, distribution, mode
|
||
|
):
|
||
|
|
||
|
if mode == "graph" and backend.is_tpu_strategy(distribution):
|
||
|
self.skipTest("partial batch not supported with TPU in graph mode.")
|
||
|
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
inputs = np.zeros((20, 3), dtype=np.float32)
|
||
|
# steps/steps_per_epoch are calculated when using numpy arrays as
|
||
|
# input data.
|
||
|
predict_with_numpy = model.predict(inputs, batch_size=10)
|
||
|
|
||
|
predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
|
||
|
inputs
|
||
|
)
|
||
|
|
||
|
self.assertEqual(
|
||
|
keras.backend.get_value(
|
||
|
tf.data.experimental.cardinality(predict_dataset)
|
||
|
),
|
||
|
tf.data.experimental.UNKNOWN_CARDINALITY,
|
||
|
)
|
||
|
|
||
|
predict_with_ds = model.predict(predict_dataset)
|
||
|
self.assertAllClose(
|
||
|
predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_on_dataset_with_unknown_cardinality_without_steps(
|
||
|
self, distribution, mode
|
||
|
):
|
||
|
# TODO(b/155867206): Investigate why this test occasionally segfaults on
|
||
|
# TPU in eager mode.
|
||
|
if mode == "eager" and backend.is_tpu_strategy(distribution):
|
||
|
self.skipTest("caused segfault with TPU in eager mode.")
|
||
|
|
||
|
if mode == "graph" and backend.is_tpu_strategy(distribution):
|
||
|
self.skipTest("partial batch not supported with TPU in graph mode.")
|
||
|
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
inputs = np.zeros((100, 3), dtype=np.float32)
|
||
|
targets = np.zeros((100, 4), dtype=np.float32)
|
||
|
# steps/steps_per_epoch are calculated when using numpy arrays as
|
||
|
# input data.
|
||
|
fit_with_numpy = model.fit(
|
||
|
inputs, targets, epochs=1, batch_size=10
|
||
|
).history
|
||
|
fit_with_numpy_multiple_epochs = model.fit(
|
||
|
inputs, targets, epochs=2, batch_size=10
|
||
|
).history
|
||
|
eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
|
||
|
predict_with_numpy = model.predict(inputs, batch_size=10)
|
||
|
|
||
|
dataset = convert_numpy_to_dataset_with_unknown_cardinality(
|
||
|
inputs, targets
|
||
|
)
|
||
|
predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
|
||
|
inputs
|
||
|
)
|
||
|
|
||
|
self.assertEqual(
|
||
|
keras.backend.get_value(
|
||
|
tf.data.experimental.cardinality(dataset)
|
||
|
),
|
||
|
tf.data.experimental.UNKNOWN_CARDINALITY,
|
||
|
)
|
||
|
self.assertEqual(
|
||
|
keras.backend.get_value(
|
||
|
tf.data.experimental.cardinality(predict_dataset)
|
||
|
),
|
||
|
tf.data.experimental.UNKNOWN_CARDINALITY,
|
||
|
)
|
||
|
|
||
|
eval_with_ds = model.evaluate(dataset)
|
||
|
predict_with_ds = model.predict(predict_dataset)
|
||
|
self.assertAllClose(
|
||
|
eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
|
||
|
fit_with_ds = model.fit(dataset, epochs=1).history
|
||
|
fit_with_ds_multiple_epochs = model.fit(dataset, epochs=2).history
|
||
|
self.assertAllClose(
|
||
|
fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
fit_with_numpy_multiple_epochs,
|
||
|
fit_with_ds_multiple_epochs,
|
||
|
atol=1e-4,
|
||
|
rtol=1e-4,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
def test_on_dataset_with_unknown_cardinality(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(
|
||
|
tf.compat.v1.train.GradientDescentOptimizer(0.001),
|
||
|
loss,
|
||
|
metrics=metrics,
|
||
|
)
|
||
|
|
||
|
inputs = np.zeros((1000, 3), dtype=np.float32)
|
||
|
targets = np.zeros((1000, 4), dtype=np.float32)
|
||
|
# steps/steps_per_epoch are calculated when using numpy arrays as
|
||
|
# input data.
|
||
|
eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
|
||
|
predict_with_numpy = model.predict(inputs, batch_size=10)
|
||
|
|
||
|
dataset = convert_numpy_to_dataset_with_unknown_cardinality(
|
||
|
inputs, targets
|
||
|
)
|
||
|
predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
|
||
|
inputs
|
||
|
)
|
||
|
|
||
|
self.assertEqual(
|
||
|
keras.backend.get_value(
|
||
|
tf.data.experimental.cardinality(dataset)
|
||
|
),
|
||
|
tf.data.experimental.UNKNOWN_CARDINALITY,
|
||
|
)
|
||
|
self.assertEqual(
|
||
|
keras.backend.get_value(
|
||
|
tf.data.experimental.cardinality(predict_dataset)
|
||
|
),
|
||
|
tf.data.experimental.UNKNOWN_CARDINALITY,
|
||
|
)
|
||
|
|
||
|
eval_with_ds = model.evaluate(dataset, steps=100)
|
||
|
predict_with_ds = model.predict(predict_dataset, steps=100)
|
||
|
self.assertAllClose(
|
||
|
eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4
|
||
|
)
|
||
|
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError, "Number of steps could not be inferred"
|
||
|
):
|
||
|
model.fit(dataset, epochs=1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_fit_eval_and_predict_methods_on_dataset(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
dataset = get_dataset(distribution)
|
||
|
|
||
|
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
|
||
|
model.evaluate(dataset, steps=2, verbose=1)
|
||
|
model.predict(get_predict_dataset(distribution), steps=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
strategy_and_optimizer_combinations()
|
||
|
)
|
||
|
def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer):
|
||
|
with self.cached_session():
|
||
|
|
||
|
with distribution.scope():
|
||
|
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer(), loss)
|
||
|
|
||
|
dataset = get_dataset(distribution)
|
||
|
|
||
|
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
|
||
|
model.evaluate(dataset, steps=2, verbose=1)
|
||
|
model.predict(get_predict_dataset(distribution), steps=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_dataset_wrong_input_shape(self, distribution, mode):
|
||
|
if mode == "graph":
|
||
|
self.skipTest(
|
||
|
"TODO(b/120943676, b/120957836): Re-enable for graph once the "
|
||
|
"validation code is restored."
|
||
|
)
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
# Wrong input shape
|
||
|
inputs = np.zeros((10, 5), dtype=np.float32)
|
||
|
targets = np.zeros((10, 4), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
|
||
|
dataset = dataset.repeat(100)
|
||
|
dataset = dataset.batch(10)
|
||
|
|
||
|
with self.assertRaisesRegex(ValueError, "is incompatible with"):
|
||
|
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu # noqa: E501
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_dataset_external_batch_input_validation(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
# Batching is done outside tf.data's `batch`
|
||
|
inputs = np.zeros((100, 10, 3), dtype=np.float32)
|
||
|
targets = np.zeros((100, 10, 4), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
|
||
|
dataset = dataset.repeat(100)
|
||
|
|
||
|
model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus, # noqa: E501
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_learning_phase_value(self, distribution):
|
||
|
# TODO(anjalisridhar): Modify this test to use Lambdas since we can
|
||
|
# compare meaningful values. Currently we don't pass the learning phase
|
||
|
# if the Lambda layer uses the learning phase.
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
x = keras.layers.Input(shape=(1,), name="input")
|
||
|
y = keras.layers.Dense(1, kernel_initializer="ones")(x)
|
||
|
z = keras.layers.Dropout(0.9999)(y)
|
||
|
model = keras.Model(x, z)
|
||
|
initial_weights = model.get_weights()
|
||
|
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.005)
|
||
|
loss = "mse"
|
||
|
metrics = ["acc"]
|
||
|
model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
batch_size = 8
|
||
|
if isinstance(
|
||
|
distribution,
|
||
|
(
|
||
|
tf.distribute.MirroredStrategy,
|
||
|
tf.compat.v1.distribute.MirroredStrategy,
|
||
|
),
|
||
|
):
|
||
|
# MirroredStrategy uses global batch size.
|
||
|
batch_size = 8 * distribution.num_replicas_in_sync
|
||
|
|
||
|
inputs = np.ones((10, 1), dtype=np.float32)
|
||
|
targets = np.ones((10, 1), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
|
||
|
dataset = dataset.repeat().batch(batch_size)
|
||
|
hist = model.fit(dataset, epochs=1, steps_per_epoch=20, verbose=1)
|
||
|
self.assertAlmostEqual(hist.history["acc"][0], 0, 0)
|
||
|
|
||
|
with distribution.scope():
|
||
|
model.set_weights(initial_weights)
|
||
|
# TODO(psv/anjalisridhar): Enable these lines after we fix
|
||
|
# b/117431185. evaluate_output = model.evaluate(dataset, steps=20)
|
||
|
# self.assertAlmostEqual(evaluate_output[1], 1, 0)
|
||
|
|
||
|
inputs = np.ones((10, 1), dtype=np.float32)
|
||
|
predict_dataset = tf.data.Dataset.from_tensor_slices(inputs)
|
||
|
|
||
|
predict_dataset = predict_dataset.repeat().batch(batch_size)
|
||
|
output = model.predict(predict_dataset, steps=10)
|
||
|
# `predict` runs for 10 steps
|
||
|
ref_output = np.ones((160, 1), dtype=np.float32)
|
||
|
self.assertArrayNear(output, ref_output, 1e-1)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def testOptimizerWithCallbacks(self, distribution):
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
model = get_model()
|
||
|
optimizer = gradient_descent_keras.SGD(0.01)
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
dataset = get_dataset(distribution)
|
||
|
|
||
|
def schedule(_):
|
||
|
return 0.001
|
||
|
|
||
|
model.fit(
|
||
|
dataset,
|
||
|
epochs=1,
|
||
|
steps_per_epoch=2,
|
||
|
verbose=0,
|
||
|
callbacks=[keras.callbacks.LearningRateScheduler(schedule)],
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
0.001, keras.backend.get_value(model.optimizer.lr)
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only(),
|
||
|
tf.__internal__.test.combinations.combine(batch_size=[4, 6]),
|
||
|
)
|
||
|
)
|
||
|
def test_evaluate_with_dataset_with_partial_batch(
|
||
|
self, distribution, batch_size
|
||
|
):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
metrics = ["mae", keras.metrics.CategoricalAccuracy()]
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = get_model()
|
||
|
model_with_ds_strategy.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
cpu_model = get_model()
|
||
|
cpu_model.compile(optimizer, loss, metrics=metrics)
|
||
|
|
||
|
x = np.random.random((10, 3)).astype("float32")
|
||
|
y = np.random.random((10, 4)).astype("float32")
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((x, y))
|
||
|
|
||
|
# As sample size is 10, we make the last batch a partial batch.
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
dataset_with_partial_batch = dataset.batch(batch_size)
|
||
|
|
||
|
# We don't compare the loss as loss is currently not computed as
|
||
|
# metric in Keras, the loss value is inaccurate for last partial
|
||
|
# batch due to more weights for the last batch samples.
|
||
|
steps = np.ceil(10.0 / batch_size)
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.evaluate(
|
||
|
dataset_with_partial_batch, steps=steps
|
||
|
)[1:],
|
||
|
cpu_model.evaluate(dataset_with_partial_batch, steps=steps)[1:],
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.evaluate(dataset_with_partial_batch)[1:],
|
||
|
cpu_model.evaluate(dataset_with_partial_batch)[1:],
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
)
|
||
|
def test_predict_with_dataset_with_partial_batch(self, distribution):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = get_model()
|
||
|
model_with_ds_strategy.compile(optimizer, loss)
|
||
|
|
||
|
cpu_model = get_model()
|
||
|
cpu_model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.random.random((10, 3)).astype(np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((inputs))
|
||
|
|
||
|
# As sample size is 10, we batch by 4 so that the last batch is
|
||
|
# a partial batch.
|
||
|
dataset_with_partial_batch = dataset.batch(4)
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.predict(
|
||
|
dataset_with_partial_batch, steps=3
|
||
|
),
|
||
|
cpu_model.predict(dataset_with_partial_batch, steps=3),
|
||
|
atol=1e-5,
|
||
|
rtol=1e-5,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
tpu_strategy_combinations_graph_only()
|
||
|
)
|
||
|
)
|
||
|
def test_predict_multi_output_model_with_dataset_with_partial_batch(
|
||
|
self, distribution
|
||
|
):
|
||
|
with self.cached_session():
|
||
|
optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
|
||
|
loss = "mse"
|
||
|
|
||
|
with distribution.scope():
|
||
|
model_with_ds_strategy = (
|
||
|
simple_multi_inputs_multi_outputs_model()
|
||
|
)
|
||
|
model_with_ds_strategy.compile(optimizer, loss)
|
||
|
|
||
|
cpu_model = simple_multi_inputs_multi_outputs_model()
|
||
|
cpu_model.compile(optimizer, loss)
|
||
|
|
||
|
input_data, _ = get_multi_inputs_multi_outputs_data()
|
||
|
input_dict = {
|
||
|
"input_a": input_data["input_a"],
|
||
|
"input_b": input_data["input_b"],
|
||
|
}
|
||
|
|
||
|
dataset = tf.data.Dataset.from_tensor_slices(input_dict)
|
||
|
|
||
|
# As sample size is 200, we batch by 18 using 12 steps per epoch so
|
||
|
# that the last batch is a partial batch.
|
||
|
dataset_with_partial_batch = dataset.batch(18)
|
||
|
cpu_model.set_weights(model_with_ds_strategy.get_weights())
|
||
|
|
||
|
self.assertAllClose(
|
||
|
model_with_ds_strategy.predict(
|
||
|
dataset_with_partial_batch, steps=12
|
||
|
),
|
||
|
cpu_model.predict(dataset_with_partial_batch, steps=12),
|
||
|
atol=1e-4,
|
||
|
rtol=1e-4,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations_minus_default()
|
||
|
)
|
||
|
def test_match_model_input_matches_with_dataset_tensors(self, distribution):
|
||
|
def _create_model_input_output_tensors():
|
||
|
input_a = keras.layers.Input(
|
||
|
shape=(16,), name="z_input_sorted_last"
|
||
|
)
|
||
|
input_b = keras.layers.Input(
|
||
|
shape=(32,), name="a_input_sorted_first"
|
||
|
)
|
||
|
intermediate_a = keras.layers.Dense(10)(input_a)
|
||
|
intermediate_b = keras.layers.Dense(10)(input_b)
|
||
|
merged = keras.layers.Add()([intermediate_a, intermediate_b])
|
||
|
output = keras.layers.Dense(2)(merged)
|
||
|
return input_a, input_b, output
|
||
|
|
||
|
input_dict = {
|
||
|
"z_input_sorted_last": np.random.rand(32, 16).astype(np.float32),
|
||
|
"a_input_sorted_first": np.random.rand(32, 32).astype(np.float32),
|
||
|
}
|
||
|
target = np.ones((32, 2), dtype=np.float32)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((input_dict, target))
|
||
|
dataset = dataset.batch(4, drop_remainder=True)
|
||
|
|
||
|
with self.cached_session():
|
||
|
with distribution.scope():
|
||
|
input_a, input_b, output = _create_model_input_output_tensors()
|
||
|
# `input_a`, which has input name that comes last in
|
||
|
# alphanumeric order, is the first input of the model input
|
||
|
# layers. If tensors from `input_dict` is blindly flattened and
|
||
|
# passed to model inputs incorrectly, this would result in
|
||
|
# `input_a` input layer matching with tensor
|
||
|
# `a_input_sorted_first` and would result in shape mismatch.
|
||
|
model_with_array_input = keras.models.Model(
|
||
|
inputs=[input_a, input_b], outputs=output
|
||
|
)
|
||
|
model_with_array_input.compile("sgd", "mse")
|
||
|
model_weights = model_with_array_input.get_weights()
|
||
|
model_with_array_input_fit = model_with_array_input.fit(
|
||
|
dataset, steps_per_epoch=1, epochs=1
|
||
|
).history
|
||
|
|
||
|
input_a, input_b, output = _create_model_input_output_tensors()
|
||
|
model_with_dict_input = keras.models.Model(
|
||
|
inputs={
|
||
|
"z_input_sorted_last": input_a,
|
||
|
"a_input_sorted_first": input_b,
|
||
|
},
|
||
|
outputs=output,
|
||
|
)
|
||
|
model_with_dict_input.compile("sgd", "mse")
|
||
|
model_with_dict_input.set_weights(model_weights)
|
||
|
model_with_dict_input_fit = model_with_dict_input.fit(
|
||
|
dataset, steps_per_epoch=1, epochs=1
|
||
|
).history
|
||
|
self.assertAllClose(
|
||
|
model_with_dict_input_fit,
|
||
|
model_with_array_input_fit,
|
||
|
atol=1e-4,
|
||
|
rtol=1e-4,
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_tpu, mode=["graph", "eager"]
|
||
|
)
|
||
|
+ tf.__internal__.test.combinations.combine(
|
||
|
distribution=multi_worker_mirrored_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_dataset_with_sample_weights(self, distribution):
|
||
|
with self.cached_session(), distribution.scope():
|
||
|
model = get_sample_weights_model()
|
||
|
optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.array([[0], [1], [2], [3]], np.float32)
|
||
|
targets = np.array([[2], [4], [6], [8]], np.float32)
|
||
|
sample_weights = np.array([0.25, 0.5, 0.75, 1], np.float32)
|
||
|
ds = tf.data.Dataset.from_tensor_slices(
|
||
|
(inputs, targets, sample_weights)
|
||
|
).batch(2)
|
||
|
result = model.evaluate(ds, verbose=1)
|
||
|
|
||
|
# The per sample loss is multiplied by the corresponding sample
|
||
|
# weight. The average of these weighted losses is the return value
|
||
|
# of the `evaluate` call. For example, in the test above the average
|
||
|
# weighted loss is calculated in the following manner:
|
||
|
# batch_1 = (((2-0)^2) * 0.25 + ((4-1)^2) * 0.5) / 2 = 5.5 / 2 =
|
||
|
# 2.75
|
||
|
# batch_2 = (((6-2)^2 * 0.75) + ((8-3)^2 * 1)) / 2 = 37 / 2 = 18.5
|
||
|
# final result = (batch_1 + batch_2) / 2 = 10.625.
|
||
|
# The first time we divide by number of input samples and the second
|
||
|
# time we divide by number of steps/batches that the loss is
|
||
|
# aggregated over.
|
||
|
self.assertAllClose(result, 10.625)
|
||
|
|
||
|
# We now test without passing sample_weights:
|
||
|
# batch_1 = ((2-0)^2) + ((4-1)^2) / 2 = 13 / 2 = 6.5
|
||
|
# batch_2 = ((6-2)^2) + ((8-3)^2) / 2 = 41 / 2 = 20.5
|
||
|
# final result = (batch_1 + batch_2) / 2 = 27 / 2 = 13.5
|
||
|
ds = tf.data.Dataset.from_tensor_slices((inputs, targets)).batch(2)
|
||
|
result = model.evaluate(ds, verbose=1)
|
||
|
self.assertAllClose(result, 13.5)
|
||
|
|
||
|
|
||
|
class TestDistributionStrategyWithDatasetsFile(
|
||
|
tf.test.TestCase, parameterized.TestCase
|
||
|
):
|
||
|
def setUp(self):
|
||
|
super().setUp()
|
||
|
self.input_file_name = os.path.join(
|
||
|
self.get_temp_dir(), "input.tfrecord"
|
||
|
)
|
||
|
inputs = np.zeros((20, 3), dtype=np.float32)
|
||
|
input_dataset = tf.data.Dataset.from_tensor_slices(inputs)
|
||
|
input_dataset = input_dataset.map(tf.io.serialize_tensor)
|
||
|
writer = tf.data.experimental.TFRecordWriter(self.input_file_name)
|
||
|
writer.write(input_dataset)
|
||
|
|
||
|
# TODO(wxinyi): add a multi-worker test for TPU
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
multi_worker_strategy_combinations_eager_only()
|
||
|
)
|
||
|
def test_predict_on_dataset_shard_options_file_multi_worker_mirrored(
|
||
|
self, distribution, mode
|
||
|
):
|
||
|
# This test is to verify if we successfully switch auto_shard_policy of
|
||
|
# a input dataset inside model.predict with MultiWorkerMirroredStrategy
|
||
|
# to AutoShardPolicy.DATA. Since there is only one input file for
|
||
|
# multiple workers, AutoShardPolicy.AUTO or AutoShardPolicy.FILE will
|
||
|
# lead to an error. However, since we switch to AutoShardPolicy.DATA in
|
||
|
# model.predict, no error is raised.
|
||
|
del mode
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
dataset = tf.data.TFRecordDataset(self.input_file_name)
|
||
|
dataset = dataset.map(lambda x: tf.io.parse_tensor(x, tf.float32))
|
||
|
|
||
|
dummy_op = lambda inp: True
|
||
|
|
||
|
dataset = dataset.filter(dummy_op).batch(8, drop_remainder=True)
|
||
|
|
||
|
options = tf.data.Options()
|
||
|
options.experimental_distribute.auto_shard_policy = (
|
||
|
tf.data.experimental.AutoShardPolicy.FILE
|
||
|
)
|
||
|
dataset = dataset.with_options(options)
|
||
|
|
||
|
model.predict(dataset, steps=1)
|
||
|
|
||
|
|
||
|
class TestRegularizerLoss(tf.test.TestCase, parameterized.TestCase):
|
||
|
class IdentityRegularizer(keras.regularizers.Regularizer):
|
||
|
def __call__(self, x):
|
||
|
return tf.identity(x)
|
||
|
|
||
|
class AddLayer(keras.layers.Layer):
|
||
|
def build(self, _):
|
||
|
self.v = self.add_weight(
|
||
|
"v",
|
||
|
(),
|
||
|
initializer="ones",
|
||
|
regularizer=TestRegularizerLoss.IdentityRegularizer(),
|
||
|
)
|
||
|
|
||
|
def call(self, inputs):
|
||
|
return inputs + self.v
|
||
|
|
||
|
@staticmethod
|
||
|
def loss_fn(_, y_pred):
|
||
|
return tf.reduce_mean(y_pred)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_combinations_minus_default()
|
||
|
)
|
||
|
)
|
||
|
def test_regularizer_loss(self, distribution):
|
||
|
batch_size = 2
|
||
|
if not distributed_training_utils.global_batch_size_supported(
|
||
|
distribution
|
||
|
):
|
||
|
batch_size //= distribution.num_replicas_in_sync
|
||
|
|
||
|
# Given an input x, which is always 1, and variable v, this model
|
||
|
# computes Loss=x+v+regularizer_loss, where regularizer_loss=v and
|
||
|
# the variable is initialized to 1. Therefore, this model computes
|
||
|
# Loss=1+2v, and so the gradient dLoss/dv = 2. This gradient of 2 is
|
||
|
# averaged over all examples in a batch and then multiplied by the
|
||
|
# learning rate of 1. As a result, the model update for one batch
|
||
|
# should subtract 2 from v, resulting in v being -1. If the
|
||
|
# regularizer loss is not scaled correctly by number of replicas,
|
||
|
# the variable value will be incorrect when number of replicas >1.
|
||
|
# For e.g. it will be -2 if num replicas = 2.
|
||
|
with distribution.scope():
|
||
|
x = keras.layers.Input(shape=(1,), batch_size=batch_size)
|
||
|
y = TestRegularizerLoss.AddLayer()(x)
|
||
|
model = keras.models.Model(inputs=x, outputs=y)
|
||
|
opt = gradient_descent_keras.SGD(1.0)
|
||
|
model.compile(opt, loss=TestRegularizerLoss.loss_fn)
|
||
|
model.fit(
|
||
|
x=np.array([[1.0], [1.0]], dtype=np.float32),
|
||
|
y=np.array([[1.0], [1.0]], dtype=np.float32),
|
||
|
batch_size=batch_size,
|
||
|
)
|
||
|
v = model.get_weights()[0]
|
||
|
self.assertEqual(-1.0, v)
|
||
|
|
||
|
|
||
|
@test_utils.run_all_without_tensor_float_32(
|
||
|
"Uses Dense layers, which call matmul"
|
||
|
)
|
||
|
class TestDistributionStrategyWithKerasModels(
|
||
|
tf.test.TestCase, parameterized.TestCase
|
||
|
):
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_distribution_strategy_on_sequential_model(self, distribution):
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = simple_sequential_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.zeros((20, 10), np.float32)
|
||
|
targets = np.zeros((20, 2), np.float32)
|
||
|
|
||
|
model.fit(inputs, targets, epochs=1, batch_size=10)
|
||
|
model.predict(inputs, batch_size=10)
|
||
|
model.evaluate(inputs, targets, batch_size=10)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_combinations()
|
||
|
)
|
||
|
def test_distribution_strategy_on_functional_model(self, distribution):
|
||
|
with distribution.scope():
|
||
|
optimizer_fn = gradient_descent_keras.SGD
|
||
|
optimizer = optimizer_fn(learning_rate=0.001)
|
||
|
model = get_model()
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
inputs = np.zeros((64, 3), dtype=np.float32)
|
||
|
targets = np.zeros((64, 4), dtype=np.float32)
|
||
|
|
||
|
model.fit(inputs, targets, epochs=1)
|
||
|
model.predict(inputs)
|
||
|
model.evaluate(inputs, targets)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_distributed_dataset(self, distribution):
|
||
|
with distribution.scope():
|
||
|
|
||
|
class CBCounter(keras.callbacks.Callback):
|
||
|
def __init__(self):
|
||
|
self.epochs = 0
|
||
|
self.train_batches = 0
|
||
|
self.test_batches = 0
|
||
|
|
||
|
def on_epoch_end(self, batch, logs=None):
|
||
|
self.epochs += 1
|
||
|
|
||
|
def on_train_batch_end(self, batch, logs=None):
|
||
|
self.train_batches += 1
|
||
|
|
||
|
def on_test_batch_end(self, batch, logs=None):
|
||
|
self.test_batches += 1
|
||
|
|
||
|
model = keras.Sequential([keras.layers.Dense(1)])
|
||
|
model.compile("sgd", "mse")
|
||
|
cb_counter = CBCounter()
|
||
|
|
||
|
x, y = np.ones((100, 10)), np.ones((100, 1))
|
||
|
ds = tf.data.Dataset.from_tensor_slices((x, y))
|
||
|
ds = ds.batch(10).repeat(2)
|
||
|
ds = distribution.experimental_distribute_dataset(ds)
|
||
|
|
||
|
val_ds = tf.data.Dataset.from_tensor_slices((x, y))
|
||
|
val_ds = val_ds.batch(20)
|
||
|
val_ds = distribution.experimental_distribute_dataset(val_ds)
|
||
|
|
||
|
model.fit(
|
||
|
ds,
|
||
|
steps_per_epoch=10,
|
||
|
validation_data=val_ds,
|
||
|
validation_steps=5,
|
||
|
epochs=2,
|
||
|
callbacks=[cb_counter],
|
||
|
)
|
||
|
|
||
|
self.assertEqual(cb_counter.train_batches, 20)
|
||
|
self.assertEqual(cb_counter.test_batches, 10)
|
||
|
self.assertEqual(cb_counter.epochs, 2)
|
||
|
|
||
|
# Check for `steps_per_epoch`.
|
||
|
if distribution.num_replicas_in_sync > 1:
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError, "distributed dataset, you must specify"
|
||
|
):
|
||
|
model.fit(ds, epochs=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_distributed_datasets_from_function(self, distribution):
|
||
|
with distribution.scope():
|
||
|
|
||
|
class CBCounter(keras.callbacks.Callback):
|
||
|
def __init__(self):
|
||
|
self.epochs = 0
|
||
|
self.train_batches = 0
|
||
|
self.test_batches = 0
|
||
|
|
||
|
def on_epoch_end(self, batch, logs=None):
|
||
|
self.epochs += 1
|
||
|
|
||
|
def on_train_batch_end(self, batch, logs=None):
|
||
|
self.train_batches += 1
|
||
|
|
||
|
def on_test_batch_end(self, batch, logs=None):
|
||
|
self.test_batches += 1
|
||
|
|
||
|
model = keras.Sequential([keras.layers.Dense(1)])
|
||
|
model.compile("sgd", "mse")
|
||
|
cb_counter = CBCounter()
|
||
|
|
||
|
def make_dataset(_):
|
||
|
x, y = np.ones((100, 10)), np.ones((100, 1))
|
||
|
ds = tf.data.Dataset.from_tensor_slices((x, y))
|
||
|
ds = ds.batch(5).repeat()
|
||
|
return ds
|
||
|
|
||
|
ds = distribution.distribute_datasets_from_function(make_dataset)
|
||
|
val_ds = distribution.distribute_datasets_from_function(
|
||
|
make_dataset
|
||
|
)
|
||
|
|
||
|
model.fit(
|
||
|
ds,
|
||
|
steps_per_epoch=10,
|
||
|
validation_data=val_ds,
|
||
|
validation_steps=5,
|
||
|
epochs=2,
|
||
|
callbacks=[cb_counter],
|
||
|
)
|
||
|
|
||
|
self.assertEqual(cb_counter.train_batches, 20)
|
||
|
self.assertEqual(cb_counter.test_batches, 10)
|
||
|
self.assertEqual(cb_counter.epochs, 2)
|
||
|
|
||
|
# Check for `steps_per_epoch`.
|
||
|
if distribution.num_replicas_in_sync > 1:
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError, "distributed dataset, you must specify"
|
||
|
):
|
||
|
model.fit(ds, epochs=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_host_training_loop(self, distribution):
|
||
|
if isinstance(distribution, tf.distribute.MultiWorkerMirroredStrategy):
|
||
|
self.skipTest("b/172032817")
|
||
|
with distribution.scope():
|
||
|
inputs = keras.Input((10, 10, 3))
|
||
|
x = keras.layers.Conv2D(3, kernel_size=3)(inputs)
|
||
|
x = keras.layers.Flatten()(x)
|
||
|
outputs = keras.layers.Dense(1)(x)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
|
||
|
model.compile("sgd", "mse", steps_per_execution=10)
|
||
|
|
||
|
bc = BatchCountingCB()
|
||
|
x, y = np.ones((100, 10, 10, 3)), np.ones((100, 1))
|
||
|
model.fit(x, y, batch_size=2, epochs=1, callbacks=[bc])
|
||
|
self.assertEqual(bc.train_begin_batches, [0, 10, 20, 30, 40])
|
||
|
self.assertEqual(bc.train_end_batches, [9, 19, 29, 39, 49])
|
||
|
|
||
|
model.evaluate(x, y, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.test_begin_batches, [0, 10, 20, 30, 40])
|
||
|
self.assertEqual(bc.test_end_batches, [9, 19, 29, 39, 49])
|
||
|
|
||
|
model.predict(x, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.predict_begin_batches, [0, 10, 20, 30, 40])
|
||
|
self.assertEqual(bc.predict_end_batches, [9, 19, 29, 39, 49])
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_host_training_loop_last_partial_execution(self, distribution):
|
||
|
if isinstance(distribution, tf.distribute.MultiWorkerMirroredStrategy):
|
||
|
self.skipTest("b/172032817")
|
||
|
with distribution.scope():
|
||
|
inputs = keras.Input(10)
|
||
|
outputs = keras.layers.Dense(1)(inputs)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
|
||
|
model.compile("sgd", "mse", steps_per_execution=20)
|
||
|
|
||
|
bc = BatchCountingCB()
|
||
|
x, y = np.ones((100, 10)), np.ones((100, 1))
|
||
|
model.fit(x, y, batch_size=2, epochs=1, callbacks=[bc])
|
||
|
self.assertEqual(bc.train_begin_batches, [0, 20, 40])
|
||
|
self.assertEqual(bc.train_end_batches, [19, 39, 49])
|
||
|
|
||
|
model.evaluate(x, y, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.test_begin_batches, [0, 20, 40])
|
||
|
self.assertEqual(bc.test_end_batches, [19, 39, 49])
|
||
|
|
||
|
model.predict(x, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.predict_begin_batches, [0, 20, 40])
|
||
|
self.assertEqual(bc.predict_end_batches, [19, 39, 49])
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_host_training_loop_dataset_unknown_size(self, distribution):
|
||
|
if isinstance(distribution, tf.distribute.MultiWorkerMirroredStrategy):
|
||
|
self.skipTest("b/172032817")
|
||
|
with distribution.scope():
|
||
|
inputs = keras.Input(10)
|
||
|
outputs = keras.layers.Dense(1)(inputs)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
|
||
|
model.compile("sgd", "mse", steps_per_execution=20)
|
||
|
|
||
|
x, y = np.ones((100, 10)), np.ones((100, 1))
|
||
|
ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(2)
|
||
|
ds = ds.filter(lambda *args, **kwargs: True) # Makes the size UNKNOWN.
|
||
|
bc = BatchCountingCB()
|
||
|
|
||
|
with self.assertRaisesRegex(ValueError, "steps_per_execution"):
|
||
|
model.fit(ds, epochs=2, callbacks=[bc])
|
||
|
|
||
|
train_ds = ds.repeat(2)
|
||
|
model.fit(train_ds, steps_per_epoch=50, epochs=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.train_begin_batches, [0, 20, 40, 0, 20, 40])
|
||
|
self.assertEqual(bc.train_end_batches, [19, 39, 49, 19, 39, 49])
|
||
|
|
||
|
with self.assertRaisesRegex(ValueError, "steps_per_execution"):
|
||
|
model.evaluate(ds, callbacks=[bc])
|
||
|
|
||
|
test_ds = ds.repeat(2)
|
||
|
model.evaluate(test_ds, steps=50, callbacks=[bc])
|
||
|
self.assertEqual(bc.test_begin_batches, [0, 20, 40])
|
||
|
self.assertEqual(bc.test_end_batches, [19, 39, 49])
|
||
|
|
||
|
predict_ds = ds.repeat(2)
|
||
|
model.predict(predict_ds, steps=50, callbacks=[bc])
|
||
|
self.assertEqual(bc.predict_begin_batches, [0, 20, 40])
|
||
|
self.assertEqual(bc.predict_end_batches, [19, 39, 49])
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_host_training_loop_truncate_to_epoch(self, distribution):
|
||
|
if isinstance(distribution, tf.distribute.MultiWorkerMirroredStrategy):
|
||
|
self.skipTest("b/172032817")
|
||
|
with distribution.scope():
|
||
|
inputs = keras.Input(10)
|
||
|
outputs = keras.layers.Dense(1)(inputs)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
|
||
|
model.compile("sgd", "mse", steps_per_execution=500)
|
||
|
|
||
|
x, y = np.ones((100, 10)), np.ones((100, 1))
|
||
|
bc = BatchCountingCB()
|
||
|
model.fit(x, y, batch_size=2, epochs=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.train_begin_batches, [0, 0])
|
||
|
self.assertEqual(bc.train_end_batches, [49, 49])
|
||
|
|
||
|
x, y = np.ones((50, 10)), np.ones((50, 1))
|
||
|
model.evaluate(x, y, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.test_begin_batches, [0])
|
||
|
self.assertEqual(bc.test_end_batches, [24])
|
||
|
|
||
|
x = np.ones((50, 10))
|
||
|
model.predict(x, batch_size=2, callbacks=[bc])
|
||
|
self.assertEqual(bc.predict_begin_batches, [0])
|
||
|
self.assertEqual(bc.predict_end_batches, [24])
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_gradient_clipping(self, distribution):
|
||
|
class MyLayer(keras.layers.Layer):
|
||
|
def build(self, _):
|
||
|
self.v1 = tf.Variable(1.0)
|
||
|
self.v2 = tf.Variable(1.0)
|
||
|
|
||
|
def call(self, x):
|
||
|
return 3 * self.v1 - 3 * self.v2
|
||
|
|
||
|
x, y = np.ones((10, 1)), np.ones((10, 1))
|
||
|
|
||
|
with distribution.scope():
|
||
|
layer = MyLayer()
|
||
|
model = keras.Sequential([layer])
|
||
|
optimizer = gradient_descent_keras.SGD(
|
||
|
1.0, clipnorm=2.0, clipvalue=2.0
|
||
|
)
|
||
|
model.compile(optimizer, "mae")
|
||
|
|
||
|
if isinstance(
|
||
|
distribution,
|
||
|
(
|
||
|
tf.distribute.experimental.CentralStorageStrategy,
|
||
|
tf.compat.v1.distribute.experimental.CentralStorageStrategy,
|
||
|
),
|
||
|
):
|
||
|
with self.assertRaisesRegex(ValueError, "not supported"):
|
||
|
model.fit(x, y, batch_size=10, epochs=1)
|
||
|
else:
|
||
|
model.fit(x, y, batch_size=10, epochs=1)
|
||
|
self.assertAllClose(self.evaluate(layer.v1), 3.0)
|
||
|
self.assertAllClose(self.evaluate(layer.v2), -1.0)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_custom_gradient_transformation(self, distribution):
|
||
|
if isinstance(
|
||
|
distribution,
|
||
|
(
|
||
|
tf.distribute.experimental.CentralStorageStrategy,
|
||
|
tf.compat.v1.distribute.experimental.CentralStorageStrategy,
|
||
|
),
|
||
|
):
|
||
|
self.skipTest("Not supported with `CentralStorageStrategy`")
|
||
|
|
||
|
class MyLayer(keras.layers.Layer):
|
||
|
def build(self, _):
|
||
|
self.v1 = tf.Variable(1.0)
|
||
|
self.v2 = tf.Variable(-1.0)
|
||
|
|
||
|
def call(self, x):
|
||
|
return x + self.v1 + self.v2
|
||
|
|
||
|
def custom_transform(grads_and_vars):
|
||
|
# Always set gradients to 1.
|
||
|
return [(tf.ones_like(g), v) for g, v in grads_and_vars]
|
||
|
|
||
|
x, y = np.ones((10, 1)), np.ones((10, 1))
|
||
|
|
||
|
with distribution.scope():
|
||
|
layer = MyLayer()
|
||
|
model = keras.Sequential([layer])
|
||
|
optimizer = gradient_descent_keras.SGD(
|
||
|
1.0, gradient_transformers=[custom_transform]
|
||
|
)
|
||
|
model.compile(optimizer, "mae")
|
||
|
|
||
|
model.fit(x, y, batch_size=10, epochs=1)
|
||
|
self.assertAllClose(self.evaluate(layer.v1), 0.0)
|
||
|
self.assertAllClose(self.evaluate(layer.v2), -2.0)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_combinations_minus_default()
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_one_dimensional(self, distribution):
|
||
|
with distribution.scope():
|
||
|
inp = keras.layers.Input(shape=(10,))
|
||
|
out = keras.layers.Dense(3, activation="softmax")(inp)
|
||
|
model = keras.Model(inputs=[inp], outputs=[out])
|
||
|
model.compile(
|
||
|
optimizer="rmsprop",
|
||
|
loss="sparse_categorical_crossentropy",
|
||
|
metrics=["sparse_categorical_accuracy"],
|
||
|
)
|
||
|
|
||
|
x = np.random.random((64, 10)).astype("float32")
|
||
|
y = np.random.randint(3, size=64)
|
||
|
|
||
|
model.fit(x, y, epochs=1, steps_per_epoch=2)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus, # noqa: E501
|
||
|
],
|
||
|
mode=["graph", "eager"],
|
||
|
reduction=[
|
||
|
losses_utils.ReductionV2.AUTO,
|
||
|
losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
|
||
|
losses_utils.ReductionV2.SUM,
|
||
|
],
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_with_loss_reduction_types(
|
||
|
self, distribution, reduction
|
||
|
):
|
||
|
np.random.seed(_RANDOM_SEED)
|
||
|
|
||
|
def _get_model():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
x2 = keras.layers.Dense(10, kernel_initializer="zeros")(x1)
|
||
|
outputs = keras.layers.Dense(1, kernel_initializer="zeros")(x2)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
return model
|
||
|
|
||
|
x = np.random.random((64, 10))
|
||
|
y = np.random.random((64, 1))
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((x, y))
|
||
|
dataset = dataset.batch(32)
|
||
|
|
||
|
model = _get_model()
|
||
|
model.compile(
|
||
|
"sgd", loss=keras.losses.MeanSquaredError(reduction=reduction)
|
||
|
)
|
||
|
history = model.fit(dataset, steps_per_epoch=2, epochs=1, shuffle=False)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _get_model()
|
||
|
ds_model.compile(
|
||
|
"sgd", loss=keras.losses.MeanSquaredError(reduction=reduction)
|
||
|
)
|
||
|
ds_history = ds_model.fit(
|
||
|
dataset, steps_per_epoch=2, epochs=1, shuffle=False
|
||
|
)
|
||
|
self.assertArrayNear(
|
||
|
history.history["loss"], ds_history.history["loss"], 1e-5
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_combinations_minus_default()
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_with_symbolic_add_loss(
|
||
|
self, mode, distribution
|
||
|
):
|
||
|
def _make_model_with_add_loss():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
x2 = keras.layers.Dense(10, kernel_initializer="zeros")(x1)
|
||
|
outputs = keras.layers.Dense(1, kernel_initializer="zeros")(x2)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
model.add_loss(tf.reduce_mean(x1))
|
||
|
model.add_loss(tf.reduce_mean(outputs))
|
||
|
return model
|
||
|
|
||
|
x = np.ones((64, 10)).astype("float32")
|
||
|
|
||
|
model = _make_model_with_add_loss()
|
||
|
model.compile("sgd")
|
||
|
history = model.fit(x, epochs=1)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _make_model_with_add_loss()
|
||
|
ds_model.compile("sgd")
|
||
|
ds_history = ds_model.fit(x, epochs=1)
|
||
|
|
||
|
self.assertAllClose(history.history, ds_history.history)
|
||
|
|
||
|
# TODO(omalleyt): Investigate flakiness and re-enable.
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
all_strategy_minus_default_and_tpu_combinations()
|
||
|
)
|
||
|
def DISABLED_test_distribution_strategy_with_callable_add_loss(
|
||
|
self, distribution
|
||
|
):
|
||
|
def _make_model():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
x2 = keras.layers.Dense(10, kernel_initializer="zeros")(x1)
|
||
|
d = keras.layers.Dense(1, kernel_initializer="zeros")
|
||
|
outputs = d(x2)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
model.add_loss(lambda: 100.0 * tf.reduce_mean(d.kernel))
|
||
|
return model
|
||
|
|
||
|
x = np.ones((64, 10)).astype("float32")
|
||
|
y = np.ones((64, 1)).astype("float32")
|
||
|
|
||
|
model = _make_model()
|
||
|
self.assertLen(model.losses, 1)
|
||
|
|
||
|
model.compile("sgd", "mse")
|
||
|
history = model.fit(x, y, steps_per_epoch=2, epochs=1)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _make_model()
|
||
|
self.assertLen(ds_model.losses, 1)
|
||
|
ds_model.compile("sgd", "mse")
|
||
|
ds_history = ds_model.fit(x, y, steps_per_epoch=2, epochs=1)
|
||
|
|
||
|
self.assertAllClose(history.history, ds_history.history)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_minus_default_and_tpu_combinations()
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_with_add_metric_in_call(self, distribution):
|
||
|
class Bias(keras.layers.Layer):
|
||
|
def build(self, input_shape):
|
||
|
self.bias = self.add_weight(
|
||
|
name="bias", initializer="zeros", shape=()
|
||
|
)
|
||
|
|
||
|
def call(self, inputs):
|
||
|
self.add_metric(
|
||
|
tf.reduce_mean(inputs), name="bias", aggregation="mean"
|
||
|
)
|
||
|
return inputs + self.bias
|
||
|
|
||
|
def _make_model_with_add_metric():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
x2 = Bias()(x1)
|
||
|
outputs = keras.layers.Dense(1, kernel_initializer="zeros")(x2)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
return model
|
||
|
|
||
|
x = np.ones((64, 10)).astype("float32")
|
||
|
y = np.ones((64, 1)).astype("float32")
|
||
|
|
||
|
model = _make_model_with_add_metric()
|
||
|
self.assertLen(model.metrics, 1)
|
||
|
|
||
|
model.compile("sgd", "mse")
|
||
|
history = model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _make_model_with_add_metric()
|
||
|
self.assertLen(ds_model.metrics, 1)
|
||
|
ds_model.compile("sgd", "mse")
|
||
|
ds_history = ds_model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
# includes stateful loss metric in eager.
|
||
|
metrics_len = 2 if tf.executing_eagerly() else 1
|
||
|
self.assertLen(ds_model.metrics, metrics_len)
|
||
|
|
||
|
self.assertAllClose(history.history, ds_history.history)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=[
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
tf.__internal__.distribute.combinations.one_device_strategy_gpu,
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu, # noqa: E501
|
||
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus, # noqa: E501
|
||
|
],
|
||
|
mode=["eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_with_add_metric_object(self, distribution):
|
||
|
class Bias(keras.layers.Layer):
|
||
|
def build(self, input_shape):
|
||
|
self.bias = self.add_weight(
|
||
|
name="bias", initializer="zeros", shape=()
|
||
|
)
|
||
|
self.mean = keras.metrics.Mean(name="mean")
|
||
|
|
||
|
def call(self, inputs):
|
||
|
self.add_metric(self.mean(inputs))
|
||
|
return inputs + self.bias
|
||
|
|
||
|
def _make_model_with_add_metric_object():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
x2 = Bias()(x1)
|
||
|
outputs = keras.layers.Dense(1, kernel_initializer="zeros")(x2)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
return model
|
||
|
|
||
|
x = np.ones((64, 10)).astype("float32")
|
||
|
y = np.ones((64, 1)).astype("float32")
|
||
|
|
||
|
model = _make_model_with_add_metric_object()
|
||
|
self.assertLen(model.metrics, 1)
|
||
|
|
||
|
model.compile("sgd", "mse")
|
||
|
history = model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _make_model_with_add_metric_object()
|
||
|
self.assertLen(ds_model.metrics, 1)
|
||
|
ds_model.compile("sgd", "mse")
|
||
|
ds_history = ds_model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
# includes stateful loss metric in eager.
|
||
|
metrics_len = 2 if tf.executing_eagerly() else 1
|
||
|
self.assertLen(ds_model.metrics, metrics_len)
|
||
|
|
||
|
self.assertAllClose(history.history, ds_history.history)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
# TODO(phillypham): Why does validation_steps > 1 not work on TPUs?
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_minus_default_and_tpu_combinations()
|
||
|
)
|
||
|
)
|
||
|
def test_distribution_strategy_with_add_metric_outside_call(
|
||
|
self, distribution
|
||
|
):
|
||
|
def _make_model_with_add_metric():
|
||
|
inputs = keras.Input((10,))
|
||
|
x1 = keras.layers.Dense(10, kernel_initializer="zeros")(inputs)
|
||
|
outputs = keras.layers.Dense(1, kernel_initializer="zeros")(x1)
|
||
|
model = keras.Model(inputs, outputs)
|
||
|
model.add_metric(
|
||
|
tf.reduce_mean(x1), name="mid_mean", aggregation="mean"
|
||
|
)
|
||
|
return model
|
||
|
|
||
|
x = np.ones((64, 10)).astype("float32")
|
||
|
y = np.ones((64, 1)).astype("float32")
|
||
|
|
||
|
model = _make_model_with_add_metric()
|
||
|
self.assertLen(model.metrics, 1)
|
||
|
|
||
|
model.compile("sgd", "mse")
|
||
|
history = model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
|
||
|
with distribution.scope():
|
||
|
ds_model = _make_model_with_add_metric()
|
||
|
self.assertLen(ds_model.metrics, 1)
|
||
|
ds_model.compile("sgd", "mse")
|
||
|
ds_history = ds_model.fit(
|
||
|
x, y, validation_data=(x, y), validation_steps=2, epochs=2
|
||
|
)
|
||
|
# includes stateful loss metric in eager.
|
||
|
metrics_len = 2 if tf.executing_eagerly() else 1
|
||
|
self.assertLen(ds_model.metrics, metrics_len)
|
||
|
|
||
|
self.assertAllClose(history.history, ds_history.history)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_tpu
|
||
|
+ multi_worker_mirrored_strategies,
|
||
|
mode=["eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_sparse_tensor_outputs(self, distribution):
|
||
|
class ToSparse(keras.layers.Layer):
|
||
|
"""Create a sparse tensor based on a given dense tensor."""
|
||
|
|
||
|
def call(self, inputs):
|
||
|
indices = tf.where(tf.not_equal(inputs, 0))
|
||
|
values = tf.gather_nd(inputs, indices)
|
||
|
shape = tf.shape(inputs, out_type="int64")
|
||
|
return tf.SparseTensor(indices, values, dense_shape=shape)
|
||
|
|
||
|
model = keras.Sequential([ToSparse()])
|
||
|
|
||
|
# Define some input data with additional padding.
|
||
|
input_data = np.array([[1, 0, 0], [2, 3, 0]])
|
||
|
output = model.predict(input_data, batch_size=2)
|
||
|
|
||
|
expected_indices = np.array([[0, 0], [1, 0], [1, 1]])
|
||
|
expected_values = np.array([1, 2, 3])
|
||
|
expected_dense_shape = np.array([2, 3])
|
||
|
|
||
|
self.assertAllEqual(output.indices, expected_indices)
|
||
|
self.assertAllEqual(output.values, expected_values)
|
||
|
self.assertAllEqual(output.dense_shape, expected_dense_shape)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_tpu
|
||
|
+ multi_worker_mirrored_strategies,
|
||
|
mode=["eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_ragged_tensor_outputs(self, distribution):
|
||
|
class ToRagged(keras.layers.Layer):
|
||
|
"""Create a ragged tensor based on a given dense tensor."""
|
||
|
|
||
|
def __init__(self, padding, ragged_rank=1, **kwargs):
|
||
|
super().__init__(**kwargs)
|
||
|
self._padding = padding
|
||
|
self._ragged_rank = ragged_rank
|
||
|
|
||
|
def call(self, inputs):
|
||
|
return tf.RaggedTensor.from_tensor(
|
||
|
inputs, padding=self._padding, ragged_rank=self._ragged_rank
|
||
|
)
|
||
|
|
||
|
model = keras.Sequential([ToRagged(padding=0)])
|
||
|
|
||
|
# Define some input data with additional padding.
|
||
|
input_data = np.array([[1, 0, 0], [2, 3, 0]])
|
||
|
output = model.predict(input_data, batch_size=2)
|
||
|
|
||
|
expected_values = [[1], [2, 3]]
|
||
|
self.assertAllEqual(expected_values, output)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=strategies_minus_default_minus_tpu
|
||
|
+ tpu_strategies
|
||
|
+ multi_worker_mirrored_strategies,
|
||
|
mode=["eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_correctness_of_add_loss_with_merge_call(self, distribution):
|
||
|
batch_size = 32
|
||
|
|
||
|
def _get_model():
|
||
|
inputs = keras.layers.Input(shape=(1,))
|
||
|
labels = keras.layers.Input(shape=(1,))
|
||
|
x = keras.layers.Dense(10, activation="relu")(inputs)
|
||
|
y = keras.layers.Dense(1)(x)
|
||
|
model = keras.models.Model([inputs, labels], y)
|
||
|
model.add_loss(keras.losses.mean_squared_error(labels, y))
|
||
|
return model
|
||
|
|
||
|
def _get_data():
|
||
|
x_train = np.random.rand(64, 1)
|
||
|
y_train = 3 * x_train
|
||
|
x_train = x_train.astype("float32")
|
||
|
y_train = y_train.astype("float32")
|
||
|
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
|
||
|
dataset = dataset.batch(batch_size)
|
||
|
return dataset
|
||
|
|
||
|
with distribution.scope():
|
||
|
model = _get_model()
|
||
|
optimizer = gradient_descent_keras.SGD(0.2)
|
||
|
|
||
|
@tf.function
|
||
|
def train_step(dist_inputs):
|
||
|
def step_fn(inputs):
|
||
|
with tf.GradientTape() as tape:
|
||
|
logits = model(inputs)
|
||
|
|
||
|
# Invoke a merge_call()
|
||
|
tf.distribute.get_replica_context().merge_call(
|
||
|
lambda d: None
|
||
|
)
|
||
|
|
||
|
# Verify that there is only one loss on the model.
|
||
|
assert len(model.losses) == 1
|
||
|
loss_from_model = (
|
||
|
tf.reduce_sum(model.losses) * 1.0 / batch_size
|
||
|
)
|
||
|
|
||
|
# Compute loss in this loop.
|
||
|
loss = keras.losses.mean_squared_error(
|
||
|
inputs[1], logits
|
||
|
)
|
||
|
loss = tf.nn.compute_average_loss(
|
||
|
loss, global_batch_size=batch_size
|
||
|
)
|
||
|
|
||
|
# Verify that the loss computed in this loop is
|
||
|
# equivalent to the loss from the model that was added
|
||
|
# via add_loss.
|
||
|
tf.compat.v1.assert_equal(loss, loss_from_model)
|
||
|
|
||
|
grads = tape.gradient(loss, model.trainable_variables)
|
||
|
optimizer.apply_gradients(
|
||
|
zip(grads, model.trainable_variables)
|
||
|
)
|
||
|
return loss
|
||
|
|
||
|
per_replica_losses = distribution.run(
|
||
|
step_fn, args=(dist_inputs,)
|
||
|
)
|
||
|
return distribution.reduce(
|
||
|
tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None
|
||
|
)
|
||
|
|
||
|
dataset = distribution.experimental_distribute_dataset(_get_data())
|
||
|
for _ in range(2):
|
||
|
for x in dataset:
|
||
|
train_step(x)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(mode=["graph", "eager"])
|
||
|
)
|
||
|
def test_unimplemented_parameter_server_strategy(self):
|
||
|
cluster_spec = multi_worker_testing_utils.create_in_process_cluster(
|
||
|
num_workers=3, num_ps=2
|
||
|
)
|
||
|
cluster_resolver = SimpleClusterResolver(
|
||
|
cluster_spec=tf.train.ClusterSpec(cluster_spec),
|
||
|
task_type="worker",
|
||
|
task_id=1,
|
||
|
num_accelerators={"GPU": 0},
|
||
|
)
|
||
|
distribution = (
|
||
|
tf.compat.v1.distribute.experimental.ParameterServerStrategy(
|
||
|
cluster_resolver
|
||
|
)
|
||
|
)
|
||
|
|
||
|
self.assertIsInstance(
|
||
|
distribution,
|
||
|
tf.compat.v1.distribute.experimental.ParameterServerStrategy,
|
||
|
)
|
||
|
|
||
|
with self.assertRaisesRegex(
|
||
|
NotImplementedError, "ParameterServerStrategy*"
|
||
|
):
|
||
|
with distribution.scope():
|
||
|
model = simple_sequential_model()
|
||
|
optimizer = tf.compat.v1.train.RMSPropOptimizer(
|
||
|
learning_rate=0.001
|
||
|
)
|
||
|
loss = "mse"
|
||
|
model.compile(optimizer, loss)
|
||
|
|
||
|
|
||
|
# Models to exercise inserting ancillary layers with add_loss and add_metric.
|
||
|
def _functional_with_add_loss_and_metric(input_shape, num_classes, l1, l2):
|
||
|
inputs = keras.Input(input_shape, name="images")
|
||
|
x = keras.layers.Conv2D(32, kernel_size=5, activation="relu")(inputs)
|
||
|
x = keras.layers.MaxPooling2D(pool_size=2)(x)
|
||
|
x = keras.layers.Conv2D(64, kernel_size=5, activation="relu")(x)
|
||
|
x = keras.layers.MaxPooling2D(pool_size=2)(x)
|
||
|
# Apply L2 regularization to embedding. Use a mix of TensorFlow ops and
|
||
|
# layers to exercise all code paths.
|
||
|
x = keras.layers.Flatten(name="embedding")(x)
|
||
|
l2_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x), -1))
|
||
|
# Apply L1 regularization to next layer.
|
||
|
x = keras.layers.Dense(1024, activation="relu", name="sparse_embedding")(x)
|
||
|
l1_loss = keras.layers.Lambda(
|
||
|
lambda x: tf.reduce_mean(tf.reduce_sum(x, -1)), name="l1_loss"
|
||
|
)(x)
|
||
|
outputs = keras.layers.Dense(num_classes, name="logits")(x)
|
||
|
model = keras.Model(inputs=inputs, outputs=outputs)
|
||
|
# Weight regularization terms.
|
||
|
model.add_loss(keras.layers.Lambda(lambda x: x * l2)(l2_loss))
|
||
|
model.add_metric(l2_loss, aggregation="mean", name="l2_loss")
|
||
|
model.add_loss(l1_loss * l1)
|
||
|
model.add_metric(l1_loss, aggregation="mean", name="l1_loss")
|
||
|
return model
|
||
|
|
||
|
|
||
|
def _sequential_with_add_loss_and_metric(input_shape, num_classes, l1, l2):
|
||
|
model = keras.Sequential(
|
||
|
[
|
||
|
keras.layers.Conv2D(
|
||
|
32, kernel_size=5, activation="relu", input_shape=input_shape
|
||
|
),
|
||
|
keras.layers.MaxPooling2D(pool_size=2),
|
||
|
keras.layers.Conv2D(64, kernel_size=5, activation="relu"),
|
||
|
keras.layers.MaxPooling2D(pool_size=2),
|
||
|
keras.layers.Flatten(name="embedding"),
|
||
|
keras.layers.Dense(
|
||
|
1024, activation="relu", name="sparse_embedding"
|
||
|
),
|
||
|
keras.layers.Dense(num_classes, name="logits"),
|
||
|
]
|
||
|
)
|
||
|
# Extract layer outputs, add regularization terms, and rescale the metric.
|
||
|
# Use a mix of TensorFlow ops and layers to exercise all code paths.
|
||
|
x = model.get_layer("sparse_embedding").get_output_at(-1)
|
||
|
l1_loss = l1 * tf.reduce_mean(tf.reduce_sum(x, -1))
|
||
|
model.add_loss(l1_loss)
|
||
|
model.add_metric(
|
||
|
keras.layers.Lambda(lambda x: tf.divide(x, l1))(l1_loss),
|
||
|
aggregation="mean",
|
||
|
name="l1_loss",
|
||
|
)
|
||
|
x = model.get_layer("embedding").get_output_at(-1)
|
||
|
l2_loss = keras.layers.Lambda(
|
||
|
lambda x: l2 * tf.reduce_mean(tf.reduce_sum(x * x, -1)), name="l2_loss"
|
||
|
)(x)
|
||
|
model.add_loss(l2_loss)
|
||
|
model.add_metric(l2_loss / l2, aggregation="mean", name="l2_loss")
|
||
|
return model
|
||
|
|
||
|
|
||
|
def _functional_with_layer_reuse(input_shape, num_classes, l1, l2):
|
||
|
base_model = keras.Sequential(
|
||
|
[
|
||
|
keras.layers.Conv2D(
|
||
|
32, kernel_size=5, activation="relu", input_shape=input_shape
|
||
|
),
|
||
|
keras.layers.MaxPooling2D(pool_size=2),
|
||
|
keras.layers.Conv2D(64, kernel_size=5, activation="relu"),
|
||
|
keras.layers.MaxPooling2D(pool_size=2),
|
||
|
keras.layers.Flatten(),
|
||
|
keras.layers.Dense(1024, activation="relu"),
|
||
|
keras.layers.Dense(num_classes, name="logits"),
|
||
|
]
|
||
|
)
|
||
|
inputs = keras.Input(input_shape, name="images")
|
||
|
logits = base_model(inputs)
|
||
|
model = keras.Model(inputs=inputs, outputs=logits)
|
||
|
# Reuse sequential layer and create new nodes.
|
||
|
zero_logits = base_model(tf.zeros_like(inputs))
|
||
|
one_logits = base_model(tf.ones_like(inputs))
|
||
|
# L2 loss.
|
||
|
l2_loss = tf.reduce_mean(tf.reduce_sum(tf.square(logits - zero_logits), -1))
|
||
|
model.add_loss(l2_loss * l2)
|
||
|
model.add_metric(l2_loss, aggregation="mean", name="l2_loss")
|
||
|
# L1 loss.
|
||
|
l1_loss = tf.reduce_mean(tf.reduce_sum(tf.abs(logits - one_logits), -1))
|
||
|
model.add_loss(l1_loss * l1)
|
||
|
model.add_metric(l1_loss, aggregation="mean", name="l1_loss")
|
||
|
return model
|
||
|
|
||
|
|
||
|
class TestDistributionStrategyWithMultipleAddLossAndMetricCalls(
|
||
|
tf.test.TestCase, parameterized.TestCase
|
||
|
):
|
||
|
"""Tests complex models with multiple add loss and metric calls."""
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.times(
|
||
|
all_strategy_combinations_minus_default(),
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
model_fn=[
|
||
|
_functional_with_add_loss_and_metric,
|
||
|
_sequential_with_add_loss_and_metric,
|
||
|
_functional_with_layer_reuse,
|
||
|
],
|
||
|
l1=[0.01],
|
||
|
l2=[0.1],
|
||
|
),
|
||
|
)
|
||
|
)
|
||
|
def test_fit_and_evaluate(self, distribution, model_fn, l1, l2):
|
||
|
# Make fake MNIST-like image data.
|
||
|
np.random.seed(_RANDOM_SEED)
|
||
|
dataset = tf.data.Dataset.from_tensor_slices(
|
||
|
(
|
||
|
np.random.uniform(size=(64, 28, 28, 1)).astype(np.float32),
|
||
|
np.random.randint(0, 10, size=(64,)),
|
||
|
)
|
||
|
)
|
||
|
dataset = dataset.shuffle(64).batch(
|
||
|
8 * distribution.num_replicas_in_sync, drop_remainder=True
|
||
|
)
|
||
|
# Make model with distribution strategy and initialize with dataset
|
||
|
# shape.
|
||
|
input_shape = tf.data.experimental.get_structure(dataset)[0].shape[1:]
|
||
|
with distribution.scope():
|
||
|
model = model_fn(input_shape, 10, l1, l2)
|
||
|
model.compile(
|
||
|
optimizer=keras.optimizers.adam_legacy.Adam(1e-4),
|
||
|
loss=keras.losses.SparseCategoricalCrossentropy(
|
||
|
from_logits=True,
|
||
|
reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
|
||
|
),
|
||
|
metrics=[
|
||
|
keras.metrics.SparseCategoricalAccuracy(),
|
||
|
keras.metrics.SparseCategoricalCrossentropy(
|
||
|
from_logits=True
|
||
|
),
|
||
|
],
|
||
|
)
|
||
|
# Non-eager training doesn't support steps_per_epoch=None.
|
||
|
for unused_epoch in range(2):
|
||
|
model.fit(dataset)
|
||
|
results = dict(zip(model.metrics_names, model.evaluate(dataset)))
|
||
|
# Sanity checks.
|
||
|
self.assertBetween(results["sparse_categorical_accuracy"], 0.02, 1.0)
|
||
|
self.assertGreater(results["l2_loss"], 0.0)
|
||
|
self.assertGreater(results["l1_loss"], 0.0)
|
||
|
# Assert correctness of the loss calculation and updating of metrics.
|
||
|
self.assertNear(
|
||
|
results["l1_loss"] * l1
|
||
|
+ results["l2_loss"] * l2
|
||
|
+ results["sparse_categorical_crossentropy"],
|
||
|
results["loss"],
|
||
|
1e-6,
|
||
|
)
|
||
|
|
||
|
|
||
|
class DeterministicModel(keras.Model):
|
||
|
"""Deterministic Model that always outputs the same initial result.
|
||
|
|
||
|
It verifies the `call` method is run inside the same distribution
|
||
|
strategy that the model was initially passed.
|
||
|
"""
|
||
|
|
||
|
def __init__(self, strategy):
|
||
|
super().__init__()
|
||
|
self.x = None
|
||
|
self.strategy = strategy
|
||
|
|
||
|
def build(self, input_shape):
|
||
|
self.x = tf.Variable(tf.ones(shape=()))
|
||
|
|
||
|
def call(self, inputs, training=None, mask=None):
|
||
|
active_strategy = tf.distribute.get_strategy()
|
||
|
if active_strategy is not self.strategy:
|
||
|
raise ValueError("Model must execute call w/ the original strategy")
|
||
|
return self.x * inputs
|
||
|
|
||
|
|
||
|
class TestModelCapturesStrategy(tf.test.TestCase, parameterized.TestCase):
|
||
|
"""Tests that model creation captures the strategy."""
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=all_strategies, mode=["eager"]
|
||
|
)
|
||
|
)
|
||
|
def test_fit_and_evaluate(self, distribution):
|
||
|
dataset = tf.data.Dataset.from_tensor_slices(
|
||
|
(tf.ones(shape=(64,)), tf.ones(shape=(64,)))
|
||
|
)
|
||
|
dataset = dataset.batch(8 * distribution.num_replicas_in_sync)
|
||
|
# Make model with distribution strategy
|
||
|
with distribution.scope():
|
||
|
model = DeterministicModel(distribution)
|
||
|
optimizer = keras.optimizers.adam_legacy.Adam(1e-4)
|
||
|
|
||
|
# Compile & evaluate the model outside of the distribution strategy
|
||
|
# scope
|
||
|
model.compile(
|
||
|
optimizer=optimizer,
|
||
|
loss=keras.losses.MeanSquaredError(),
|
||
|
metrics=["binary_accuracy"],
|
||
|
)
|
||
|
|
||
|
# Call `optimizer.iterations` out of strategy scope.
|
||
|
self.assertEqual(model.optimizer.iterations.numpy(), 0)
|
||
|
|
||
|
# Non-eager training doesn't support steps_per_epoch=None.
|
||
|
for unused_epoch in range(2):
|
||
|
model.fit(dataset)
|
||
|
|
||
|
results = model.evaluate(dataset)
|
||
|
results = dict(zip(model.metrics_names, results))
|
||
|
|
||
|
# Check that the metrics have a result we expect
|
||
|
self.assertEqual(results["binary_accuracy"], 1.0)
|
||
|
self.assertAllClose(results["loss"], 0.0)
|
||
|
|
||
|
# Assert that all metric/optimizer/model variables were made in the
|
||
|
# distribution strategy (Test that compile uses the captured
|
||
|
# distribution strategy)
|
||
|
metric_vars = tf.nest.flatten(
|
||
|
[metric.variables for metric in model.metrics]
|
||
|
)
|
||
|
for var in metric_vars:
|
||
|
self.assertTrue(
|
||
|
distribution.extended.variable_created_in_scope(var)
|
||
|
)
|
||
|
for var in model.optimizer._weights:
|
||
|
self.assertTrue(
|
||
|
distribution.extended.variable_created_in_scope(var)
|
||
|
)
|
||
|
for var in model.variables:
|
||
|
self.assertTrue(
|
||
|
distribution.extended.variable_created_in_scope(var)
|
||
|
)
|
||
|
|
||
|
# Make sure the metric must be created in the same scope as the model:
|
||
|
# This shouldn't raise any validation errors
|
||
|
with distribution.scope():
|
||
|
metric = keras.metrics.BinaryAccuracy()
|
||
|
model.compile(
|
||
|
optimizer=optimizer,
|
||
|
loss=keras.losses.MeanSquaredError(),
|
||
|
metrics=[metric],
|
||
|
)
|
||
|
|
||
|
# This should raise an error because the metric is constructed
|
||
|
# outside of the scope, and not by compile
|
||
|
if tf.distribute.has_strategy():
|
||
|
with self.assertRaisesRegex(
|
||
|
ValueError, "All metrics must be created in"
|
||
|
):
|
||
|
model.compile(
|
||
|
optimizer=keras.optimizers.adam_v2.Adam(1e-4),
|
||
|
loss=keras.losses.MeanSquaredError(),
|
||
|
metrics=[keras.metrics.BinaryAccuracy()],
|
||
|
)
|
||
|
|
||
|
@tf.__internal__.distribute.combinations.generate(
|
||
|
tf.__internal__.test.combinations.combine(
|
||
|
distribution=tf.__internal__.distribute.combinations.mirrored_strategy_with_one_cpu, # noqa: E501
|
||
|
mode=["eager"],
|
||
|
)
|
||
|
)
|
||
|
def test_optimizer(self, distribution):
|
||
|
temp_dir = os.path.join(self.get_temp_dir(), "ckpt")
|
||
|
|
||
|
def create_model():
|
||
|
model = keras.models.Sequential(
|
||
|
[
|
||
|
keras.layers.Dense(1),
|
||
|
]
|
||
|
)
|
||
|
model.compile(optimizer=keras.optimizers.Adam(), loss="mse")
|
||
|
model.build([None, 1]) # create weights.
|
||
|
return model
|
||
|
|
||
|
model = create_model()
|
||
|
x = y = tf.ones(shape=(1, 1))
|
||
|
model.fit(x=x, y=y, batch_size=1)
|
||
|
model.save_weights(temp_dir)
|
||
|
|
||
|
with distribution.scope():
|
||
|
model = create_model()
|
||
|
model.load_weights(temp_dir)
|
||
|
if isinstance(model.optimizer, optimizer_base.Optimizer):
|
||
|
model.optimizer.build(model.trainable_variables)
|
||
|
variables = model.optimizer.variables
|
||
|
else:
|
||
|
variables = model.optimizer.variables()
|
||
|
self.assertNotEmpty(variables)
|
||
|
self.assertTrue(
|
||
|
distributed_training_utils.is_distributed_variable(variables[0])
|
||
|
)
|
||
|
|
||
|
with distribution.scope():
|
||
|
model = create_model()
|
||
|
# create/restore slot variables outside of scope is fine.
|
||
|
model.load_weights(temp_dir)
|
||
|
if isinstance(model.optimizer, optimizer_base.Optimizer):
|
||
|
# V3 optimizer has to restore variables in scope.
|
||
|
return
|
||
|
# From this point on, the optimizer must be a V2 optimizer.
|
||
|
self.assertNotEmpty(model.optimizer.variables())
|
||
|
self.assertTrue(
|
||
|
distributed_training_utils.is_distributed_variable(
|
||
|
model.optimizer.variables()[0]
|
||
|
)
|
||
|
)
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
base_layer_utils.enable_v2_dtype_behavior()
|
||
|
tf.__internal__.distribute.multi_process_runner.test_main()
|