288 lines
10 KiB
Python
288 lines
10 KiB
Python
|
# Copyright 2019 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Base class for testing saving/loading with DS."""
|
||
|
|
|
||
|
|
import os
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import tensorflow.compat.v2 as tf
|
||
|
|
from absl.testing import parameterized
|
||
|
|
|
||
|
|
from keras.distribute import model_combinations
|
||
|
|
|
||
|
|
_RANDOM_SEED = 1337
|
||
|
|
_DEFAULT_FUNCTION_KEY = "serving_default"
|
||
|
|
|
||
|
|
_TOLERANCE = 1e-30
|
||
|
|
# TPU uses bfloat16 for computation in hardware underlying, so it has less
|
||
|
|
# precision than CPU/GPU.
|
||
|
|
_TPU_TOLERANCE = 1e-7
|
||
|
|
|
||
|
|
PREDICT_STEPS = 1
|
||
|
|
|
||
|
|
simple_models = [
|
||
|
|
model_combinations.simple_functional_model,
|
||
|
|
model_combinations.simple_sequential_model,
|
||
|
|
model_combinations.simple_subclass_model,
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
strategies = [
|
||
|
|
tf.__internal__.distribute.combinations.default_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.one_device_strategy_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_one_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_one_gpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
|
||
|
|
tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
|
||
|
|
tf.__internal__.distribute.combinations.tpu_strategy,
|
||
|
|
tf.__internal__.distribute.combinations.tpu_strategy_packed_var,
|
||
|
|
tf.__internal__.distribute.combinations.central_storage_strategy_with_two_gpus, # noqa: E501
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def simple_models_with_strategies():
|
||
|
|
return tf.__internal__.test.combinations.combine(
|
||
|
|
model_and_input=simple_models, distribution=strategies, mode=["eager"]
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def simple_models_with_strategy_pairs():
|
||
|
|
return tf.__internal__.test.combinations.combine(
|
||
|
|
model_and_input=simple_models,
|
||
|
|
distribution_for_saving=strategies,
|
||
|
|
distribution_for_restoring=strategies,
|
||
|
|
mode=["eager"],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def tfmodule_models_with_strategies():
|
||
|
|
return tf.__internal__.test.combinations.combine(
|
||
|
|
model_and_input=[model_combinations.simple_tfmodule_model],
|
||
|
|
distribution=strategies,
|
||
|
|
mode=["eager"],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def tfmodule_models_with_strategy_pairs():
|
||
|
|
return tf.__internal__.test.combinations.combine(
|
||
|
|
model_and_input=[model_combinations.simple_tfmodule_model],
|
||
|
|
distribution_for_saving=strategies,
|
||
|
|
distribution_for_restoring=strategies,
|
||
|
|
mode=["eager"],
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load_and_run_with_saved_model_api(
|
||
|
|
distribution, saved_dir, predict_dataset, output_name
|
||
|
|
):
|
||
|
|
"""Loads a saved_model using tf.saved_model API, and runs it."""
|
||
|
|
func = tf.saved_model.load(saved_dir)
|
||
|
|
if distribution:
|
||
|
|
dist_predict_dataset = distribution.experimental_distribute_dataset(
|
||
|
|
predict_dataset
|
||
|
|
)
|
||
|
|
per_replica_predict_data = next(iter(dist_predict_dataset))
|
||
|
|
result = distribution.run(
|
||
|
|
func.signatures[_DEFAULT_FUNCTION_KEY],
|
||
|
|
args=(per_replica_predict_data,),
|
||
|
|
)
|
||
|
|
result = result[output_name]
|
||
|
|
|
||
|
|
# Convert the per_replica value to a list, then concatenate them
|
||
|
|
reduced = distribution.experimental_local_results(result)
|
||
|
|
concat = tf.concat(reduced, 0)
|
||
|
|
return concat
|
||
|
|
else:
|
||
|
|
result = func.signatures[_DEFAULT_FUNCTION_KEY](
|
||
|
|
next(iter(predict_dataset))
|
||
|
|
)
|
||
|
|
return result[output_name]
|
||
|
|
|
||
|
|
|
||
|
|
class TestSavedModelBase(tf.test.TestCase, parameterized.TestCase):
|
||
|
|
"""Base class for testing saving/loading with DS."""
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
np.random.seed(_RANDOM_SEED)
|
||
|
|
tf.compat.v1.set_random_seed(_RANDOM_SEED)
|
||
|
|
self._root_dir = "base"
|
||
|
|
super().setUp()
|
||
|
|
|
||
|
|
def _save_model(self, model, saved_dir):
|
||
|
|
"""Save the given model to the given saved_dir.
|
||
|
|
|
||
|
|
This method needs to be implemented by the subclasses.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: a keras model object to save.
|
||
|
|
saved_dir: a string representing the path to save the keras model
|
||
|
|
"""
|
||
|
|
raise NotImplementedError("must be implemented in descendants")
|
||
|
|
|
||
|
|
def _load_and_run_model(
|
||
|
|
self, distribution, saved_dir, predict_dataset, output_name="output_1"
|
||
|
|
):
|
||
|
|
"""Load the model and run 1 step of predict with it.
|
||
|
|
|
||
|
|
This method must be implemented by the subclasses.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
distribution: the distribution strategy used to load the model. None
|
||
|
|
if no distribution strategy is used
|
||
|
|
saved_dir: the string representing the path where the model is saved.
|
||
|
|
predict_dataset: the data used to do the predict on the model for
|
||
|
|
cross_replica context.
|
||
|
|
output_name: the string representing the name of the output layer of
|
||
|
|
the model.
|
||
|
|
"""
|
||
|
|
|
||
|
|
raise NotImplementedError("must be implemented in descendants")
|
||
|
|
|
||
|
|
def _train_model(self, model, x_train, y_train, batch_size):
|
||
|
|
training_dataset = tf.data.Dataset.from_tensor_slices(
|
||
|
|
(x_train, y_train)
|
||
|
|
)
|
||
|
|
training_dataset = training_dataset.repeat()
|
||
|
|
training_dataset = training_dataset.batch(batch_size)
|
||
|
|
|
||
|
|
# Train the model for 1 epoch
|
||
|
|
model.fit(x=training_dataset, epochs=1, steps_per_epoch=100)
|
||
|
|
|
||
|
|
def _predict_with_model(self, distribution, model, predict_dataset):
|
||
|
|
return model.predict(predict_dataset, steps=PREDICT_STEPS)
|
||
|
|
|
||
|
|
def _get_predict_dataset(self, x_predict, batch_size):
|
||
|
|
predict_dataset = tf.data.Dataset.from_tensor_slices(x_predict)
|
||
|
|
predict_dataset = predict_dataset.repeat()
|
||
|
|
predict_dataset = predict_dataset.batch(batch_size)
|
||
|
|
return predict_dataset
|
||
|
|
|
||
|
|
def run_test_save_no_strategy_restore_strategy(
|
||
|
|
self, model_and_input, distribution
|
||
|
|
):
|
||
|
|
"""Save a model without DS, and restore it with DS."""
|
||
|
|
|
||
|
|
saved_dir = os.path.join(self.get_temp_dir(), "0")
|
||
|
|
|
||
|
|
model = model_and_input.get_model()
|
||
|
|
x_train, y_train, x_predict = model_and_input.get_data()
|
||
|
|
batch_size = model_and_input.get_batch_size()
|
||
|
|
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
|
||
|
|
|
||
|
|
self._train_model(model, x_train, y_train, batch_size)
|
||
|
|
result_before_save = self._predict_with_model(
|
||
|
|
None, model, predict_dataset
|
||
|
|
)
|
||
|
|
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
|
||
|
|
with distribution.scope():
|
||
|
|
result_after_save = self._load_and_run_model(
|
||
|
|
distribution=distribution,
|
||
|
|
saved_dir=saved_dir,
|
||
|
|
predict_dataset=predict_dataset,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertAllClose(result_before_save, result_after_save)
|
||
|
|
|
||
|
|
def run_test_save_strategy_restore_no_strategy(
|
||
|
|
self, model_and_input, distribution, save_in_scope
|
||
|
|
):
|
||
|
|
"""Save a model with DS, and restore it without DS."""
|
||
|
|
|
||
|
|
saved_dir = os.path.join(self.get_temp_dir(), "1")
|
||
|
|
|
||
|
|
with distribution.scope():
|
||
|
|
model = model_and_input.get_model()
|
||
|
|
x_train, y_train, x_predict = model_and_input.get_data()
|
||
|
|
batch_size = model_and_input.get_batch_size()
|
||
|
|
|
||
|
|
self._train_model(model, x_train, y_train, batch_size)
|
||
|
|
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
|
||
|
|
result_before_save = self._predict_with_model(
|
||
|
|
distribution, model, predict_dataset
|
||
|
|
)
|
||
|
|
|
||
|
|
if save_in_scope:
|
||
|
|
with distribution.scope():
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
else:
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
|
||
|
|
load_result = self._load_and_run_model(
|
||
|
|
distribution=None,
|
||
|
|
saved_dir=saved_dir,
|
||
|
|
predict_dataset=predict_dataset,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertAllClose(result_before_save, load_result)
|
||
|
|
|
||
|
|
def run_test_save_strategy_restore_strategy(
|
||
|
|
self,
|
||
|
|
model_and_input,
|
||
|
|
distribution_for_saving,
|
||
|
|
distribution_for_restoring,
|
||
|
|
save_in_scope,
|
||
|
|
):
|
||
|
|
"""Save a model with DS, and restore it with potentially different
|
||
|
|
DS."""
|
||
|
|
saved_dir = os.path.join(self.get_temp_dir(), "2")
|
||
|
|
|
||
|
|
with distribution_for_saving.scope():
|
||
|
|
model = model_and_input.get_model()
|
||
|
|
x_train, y_train, x_predict = model_and_input.get_data()
|
||
|
|
batch_size = model_and_input.get_batch_size()
|
||
|
|
|
||
|
|
self._train_model(model, x_train, y_train, batch_size)
|
||
|
|
predict_dataset = self._get_predict_dataset(x_predict, batch_size)
|
||
|
|
result_before_save = self._predict_with_model(
|
||
|
|
distribution_for_saving, model, predict_dataset
|
||
|
|
)
|
||
|
|
|
||
|
|
if save_in_scope:
|
||
|
|
with distribution_for_saving.scope():
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
else:
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
|
||
|
|
with distribution_for_restoring.scope():
|
||
|
|
|
||
|
|
load_result = self._load_and_run_model(
|
||
|
|
distribution=distribution_for_restoring,
|
||
|
|
saved_dir=saved_dir,
|
||
|
|
predict_dataset=predict_dataset,
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertAllClose(result_before_save, load_result)
|
||
|
|
|
||
|
|
def run_test_save_strategy(
|
||
|
|
self, model_and_input, distribution, save_in_scope
|
||
|
|
):
|
||
|
|
"""Save a model with DS."""
|
||
|
|
saved_dir = os.path.join(self.get_temp_dir(), "3")
|
||
|
|
with distribution.scope():
|
||
|
|
model = model_and_input.get_model()
|
||
|
|
x_train, y_train, _ = model_and_input.get_data()
|
||
|
|
batch_size = model_and_input.get_batch_size()
|
||
|
|
self._train_model(model, x_train, y_train, batch_size)
|
||
|
|
|
||
|
|
if save_in_scope:
|
||
|
|
with distribution.scope():
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
else:
|
||
|
|
self._save_model(model, saved_dir)
|
||
|
|
return saved_dir
|