# 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. # ============================================================================= """Test async checkpointing.""" import os import numpy as np from tensorflow.core.framework import summary_pb2 from tensorflow.python.compat import v2_compat from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.lib.io import file_io from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics as metrics_lib from tensorflow.python.ops import variable_scope from tensorflow.python.ops.losses import losses from tensorflow.python.platform import flags from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging as logging from tensorflow.python.saved_model.pywrap_saved_model import metrics from tensorflow.python.tpu import async_checkpoint from tensorflow.python.tpu import tpu_optimizer from tensorflow.python.training import basic_session_run_hooks from tensorflow.python.training import training from tensorflow_estimator.python.estimator import estimator as estimator_lib from tensorflow_estimator.python.estimator import model_fn as model_fn_lib from tensorflow_estimator.python.estimator.tpu import tpu_config from tensorflow_estimator.python.estimator.tpu import tpu_estimator FLAGS = flags.FLAGS flags.DEFINE_string('tpu', '', 'TPU to use in this test.') flags.DEFINE_string('zone', None, 'Name of GCP zone with TPU.') flags.DEFINE_string('project', None, 'Name of GCP project with TPU.') flags.DEFINE_string( 'model_dir', os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR'), 'GCS path to store model and checkpoints.') def _get_checkpoint_metrics_counts() -> (int, int): """Get the count for recorded sync and async checkpoint write durations.""" def get_count(method): proto_bytes = method(api_label=async_checkpoint._ASYNC_CHECKPOINT_V1) histogram_proto = summary_pb2.HistogramProto() histogram_proto.ParseFromString(proto_bytes) return int(histogram_proto.num) return get_count(metrics.GetCheckpointWriteDurations), get_count( metrics.GetAsyncCheckpointWriteDurations) def input_fn(params): """Return a dataset of source and target sequences for training.""" return (constant_op.constant( np.random.randn(params['batch_size'], 1000), dtype=dtypes.float32), constant_op.constant( np.random.randint(0, 10, params['batch_size']), dtype=dtypes.int32)) def model_fn(features, labels, mode, params): del params # unused with variable_scope.variable_scope('m', reuse=variable_scope.AUTO_REUSE): w = variable_scope.get_variable('W', shape=[1000, 10]) logits = math_ops.matmul(features, w) loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) if mode == model_fn_lib.ModeKeys.TRAIN: optimizer = training.RMSPropOptimizer(learning_rate=0.01) optimizer = tpu_optimizer.CrossShardOptimizer(optimizer) train_op = optimizer.minimize(loss, training.get_global_step()) return tpu_estimator.TPUEstimatorSpec( mode=model_fn_lib.ModeKeys.TRAIN, loss=loss, train_op=train_op, ) elif mode == model_fn_lib.ModeKeys.EVAL: def metric_fn(labels, logits): labels = math_ops.cast(labels, dtypes.int64) logging.info('LABELS %s %s', labels, logits) return { 'recall@1': metrics_lib.recall_at_k(labels, logits, 1), 'recall@5': metrics_lib.recall_at_k(labels, logits, 5), } loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) eval_metrics = (metric_fn, [labels, logits]) return tpu_estimator.TPUEstimatorSpec( mode=model_fn_lib.ModeKeys.EVAL, loss=loss, eval_metrics=eval_metrics) class AsyncCheckpointingTest(test.TestCase): def testAsyncCheckpointHookEnabled(self): resolver = tpu_cluster_resolver.TPUClusterResolver( tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project) checkpoint_interval = 5 config = tpu_config.RunConfig( master=resolver.master(), model_dir=os.path.join(FLAGS.model_dir, 'runconfig'), save_checkpoints_steps=1000, keep_checkpoint_max=11, # off by one tpu_config=tpu_config.TPUConfig( iterations_per_loop=checkpoint_interval,)) estimator = tpu_estimator.TPUEstimator( use_tpu=True, model_fn=model_fn, config=config, train_batch_size=32, eval_batch_size=32, predict_batch_size=1, params={}, ) max_steps = 100 mock_listener = test.mock.create_autospec( basic_session_run_hooks.CheckpointSaverListener) estimator.train( input_fn=input_fn, max_steps=max_steps, hooks=[ async_checkpoint.AsyncCheckpointSaverHook( FLAGS.model_dir, save_steps=checkpoint_interval, listeners=[mock_listener]) ]) current_step = estimator_lib._load_global_step_from_checkpoint_dir( FLAGS.model_dir) # pylint: disable=protected-access # TODO(power) -- identify a better way to count the number of checkpoints. checkpoints = file_io.get_matching_files( FLAGS.model_dir + '/model.ckpt*.meta') checkpoint_count = len(checkpoints) logging.info('Found %d checkpoints: %s', checkpoint_count, checkpoints) self.assertLessEqual(checkpoint_count, 10) self.assertEqual(current_step, max_steps) mock_listener.before_save.assert_called() mock_listener.after_save.assert_called() # save called by hook in `after_create_session` and every `after_run` num_save_calls = 1 + max_steps // checkpoint_interval sync_count_1, async_count_1 = _get_checkpoint_metrics_counts() # save might be called one extra time in `end` hook based on timing of # `_last_checkpoint_step` update in the final `after_run` call self.assertIn(sync_count_1, [num_save_calls, num_save_calls + 1]) self.assertLessEqual(async_count_1, num_save_calls) training_time_saved = metrics.GetTrainingTimeSaved( api_label=async_checkpoint._ASYNC_CHECKPOINT_V1) self.assertGreater(training_time_saved, 0) def testAsyncCheckpointHookWithoutListeners(self): resolver = tpu_cluster_resolver.TPUClusterResolver( tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project) checkpoint_interval = 5 keep_checkpoint_max = 10 config = tpu_config.RunConfig( master=resolver.master(), model_dir=os.path.join(FLAGS.model_dir, 'runconfig'), save_checkpoints_steps=1000, keep_checkpoint_max=keep_checkpoint_max+1, # off by one tpu_config=tpu_config.TPUConfig( iterations_per_loop=checkpoint_interval,)) estimator = tpu_estimator.TPUEstimator( use_tpu=True, model_fn=model_fn, config=config, train_batch_size=32, eval_batch_size=32, predict_batch_size=1, params={}, ) max_steps = 100 estimator.train( input_fn=input_fn, max_steps=max_steps, hooks=[ async_checkpoint.AsyncCheckpointSaverHook( FLAGS.model_dir, save_steps=checkpoint_interval) ]) current_step = estimator_lib._load_global_step_from_checkpoint_dir( FLAGS.model_dir) # pylint: disable=protected-access # TODO(power) -- identify a better way to count the number of checkpoints. checkpoints = file_io.get_matching_files( FLAGS.model_dir + '/model.ckpt*.meta') checkpoint_count = len(checkpoints) logging.info('Found %d checkpoints: %s', checkpoint_count, checkpoints) self.assertLessEqual(checkpoint_count, keep_checkpoint_max) self.assertEqual(current_step, max_steps) # save called by hook in `after_create_session` and every `after_run` num_save_calls = 1 + max_steps // checkpoint_interval sync_count_1, async_count_1 = _get_checkpoint_metrics_counts() # save might be called one extra time in `end` hook based on timing of # `_last_checkpoint_step` update in the final `after_run` call self.assertIn(sync_count_1, [num_save_calls, num_save_calls + 1]) self.assertLessEqual(async_count_1, num_save_calls) training_time_saved = metrics.GetTrainingTimeSaved( api_label=async_checkpoint._ASYNC_CHECKPOINT_V1) self.assertGreater(training_time_saved, 0) if __name__ == '__main__': v2_compat.disable_v2_behavior() test.main()