def test_save_load_trackable(self, distribution, optimizer): # TODO(b/123533246): Enable the test for TPU once bug is fixed if (isinstance( distribution, ( tf.distribute.experimental.TPUStrategy, tf.compat.v1.distribute.experimental.TPUStrategy, ), ) and distribution.extended.steps_per_run > 1): self.skipTest( "MultiStep TPU Strategy deadlocks with optimizer restore.") with self.cached_session(): dataset = keras_test_lib.get_dataset(distribution) with distribution.scope(): model = keras_test_lib.get_model() model.compile(optimizer(), "mse") model.fit(dataset, epochs=1, steps_per_epoch=1) weights_file = tempfile.mktemp() model.save_weights(weights_file) model_2 = keras_test_lib.get_model() model_2.compile(optimizer(), "mse") model_2.load_weights(weights_file) model_2.predict( keras_test_lib.get_predict_dataset(distribution), steps=2) model_2.fit(dataset, epochs=1, steps_per_epoch=1)
def testV2SummaryWithKerasSubclassedModel(self): # Histogram summaries require the MLIR bridge; see b/178826597#comment107. # TODO(https://github.com/tensorflow/tensorboard/issues/2885): remove this # if histogram summaries are supported fully on non-MLIR bridge or # non-MLIR bridge is no longer run. enable_histograms = test_util.is_mlir_bridge_enabled() strategy = get_tpu_strategy() with strategy.scope(): model = CustomModel(enable_histograms=enable_histograms) model.compile('sgd', 'mse') dataset = distribute_strategy_test.get_dataset(strategy) tensorboard_callback = callbacks.TensorBoard( self.summary_dir, update_freq=2) model.fit( dataset, steps_per_epoch=10, epochs=1, callbacks=[tensorboard_callback]) event_files = tf.io.gfile.glob( os.path.join(self.summary_dir, 'train', 'event*')) # Since total of 10 steps are ran and summary ops should be invoked # every 2 batches, we should see total of 5 event logs for each summary. expected_event_counts = { ('custom_model/layer_for_scalar_summary/' 'custom_scalar_summary_v2'): 5, ('custom_model/layer_for_histogram_summary/' 'custom_histogram_summary_v2'): 5 if enable_histograms else 0, } self.validate_recorded_sumary_file(event_files, expected_event_counts)
def testV2SummaryWithKerasSubclassedModel(self): strategy = get_tpu_strategy() with strategy.scope(): model = CustomModel() model.compile('sgd', 'mse') dataset = distribute_strategy_test.get_dataset(strategy) tensorboard_callback = callbacks.TensorBoard( self.summary_dir, update_freq=2) model.fit( dataset, steps_per_epoch=10, epochs=1, callbacks=[tensorboard_callback]) event_files = tf.io.gfile.glob( os.path.join(self.summary_dir, 'train', 'event*')) events_count_dictionary = { ('custom_model/layer_for_scalar_summary/' 'custom_scalar_summary_v2'): 0, ('custom_model/layer_for_histogram_summary/' 'custom_histogram_summary_v2'): 0 } # Since total of 10 steps are ran and summary ops should be invoked # every 2 batches, we should see total of 5 event logs. self.validate_recorded_sumary_file(event_files, events_count_dictionary, 5)
def test_callbacks_in_fit(self, distribution): with distribution.scope(): model = keras_test_lib.get_model() model.compile(optimizer="sgd", loss="mse", metrics=["mae"]) dataset = keras_test_lib.get_dataset(distribution) counter = Counter() epochs = 2 steps_per_epoch = 5 validation_steps = 3 model.fit( dataset, epochs=epochs, steps_per_epoch=steps_per_epoch, verbose=0, validation_data=dataset, validation_steps=validation_steps, callbacks=[counter], ) if (isinstance(distribution, tf.compat.v1.distribute.experimental.TPUStrategy) and not tf.executing_eagerly()): # TPU Strategy can have multi step training, from # extended.steps_per_run if steps_per_run = 1, then # num_batch_call_per_epoch = steps_per_epoch steps_per_run = distribution.extended.steps_per_run num_batch_call_per_epoch = steps_per_epoch // steps_per_run if steps_per_epoch % steps_per_run: num_batch_call_per_epoch += 1 else: num_batch_call_per_epoch = steps_per_epoch self.assertDictEqual( counter.method_counts, { "on_batch_begin": epochs * num_batch_call_per_epoch, "on_batch_end": epochs * num_batch_call_per_epoch, "on_epoch_begin": epochs, "on_epoch_end": epochs, "on_test_batch_begin": epochs * validation_steps, "on_test_batch_end": epochs * validation_steps, "on_test_begin": epochs, "on_test_end": epochs, "on_train_batch_begin": epochs * num_batch_call_per_epoch, "on_train_batch_end": epochs * num_batch_call_per_epoch, "on_train_begin": 1, "on_train_end": 1, }, )
def test_save_load_h5(self, distribution, optimizer): with self.cached_session(): dataset = keras_test_lib.get_dataset(distribution) with distribution.scope(): model = keras_test_lib.get_model() model.compile(optimizer(), 'mse') model.fit(dataset, epochs=1, steps_per_epoch=1) weights_file = tempfile.mktemp('.h5') model.save_weights(weights_file) model_2 = keras_test_lib.get_model() model_2.compile(optimizer(), 'mse') model_2.load_weights(weights_file) model_2.predict( keras_test_lib.get_predict_dataset(distribution), steps=2) model_2.fit(dataset, epochs=1, steps_per_epoch=1)
def test_callbacks_in_eval(self, distribution): with distribution.scope(): model = keras_test_lib.get_model() model.compile(optimizer='sgd', loss='mse', metrics=['mae']) dataset = keras_test_lib.get_dataset(distribution) counter = Counter() model.evaluate(dataset, steps=5, callbacks=[counter]) self.assertDictEqual( counter.method_counts, { 'on_test_batch_begin': 5, 'on_test_batch_end': 5, 'on_test_begin': 1, 'on_test_end': 1 })
def test_callbacks_in_eval(self, distribution): with distribution.scope(): model = keras_test_lib.get_model() model.compile(optimizer="sgd", loss="mse", metrics=["mae"]) dataset = keras_test_lib.get_dataset(distribution) counter = Counter() model.evaluate(dataset, steps=5, callbacks=[counter]) self.assertDictEqual( counter.method_counts, { "on_test_batch_begin": 5, "on_test_batch_end": 5, "on_test_begin": 1, "on_test_end": 1, }, )
def test_unsupported_features(self, distribution, mode): with self.cached_session(): with distribution.scope(): model = keras_test_lib.get_model() optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001) loss = "mse" metrics = ["mae"] model.compile(optimizer, loss, metrics=metrics) dataset = keras_test_lib.get_dataset(distribution) # Test with validation split with self.assertRaises(ValueError): model.fit( dataset, epochs=1, steps_per_epoch=2, verbose=0, validation_split=0.5, validation_steps=2, ) # Test with sample weight. sample_weight = np.random.random((10, )) with self.assertRaises(ValueError): model.fit( dataset, epochs=1, steps_per_epoch=2, verbose=0, sample_weight=sample_weight, ) # Test with not specifying the `steps` argument for dataset with infinite # cardinality. dataset = dataset.repeat() with self.assertRaises(ValueError): model.fit(dataset, epochs=1, verbose=0) with self.assertRaises(ValueError): model.evaluate(dataset, verbose=0) with self.assertRaises(ValueError): model.predict(dataset, verbose=0)
def testSummaryWithCustomTrainingLoop(self): strategy = get_tpu_strategy() writer = tf.summary.create_file_writer(self.summary_dir) with strategy.scope(): model = distribute_strategy_test.get_model() model.compile('sgd', 'mse') @tf.function def custom_function(dataset): def _custom_step(features, labels): del labels logits = model(features) with tf.summary.record_if(True), writer.as_default(): scalar_summary_v2.scalar( 'logits', tf.reduce_sum(logits), step=model.optimizer.iterations) return logits iterator = iter(dataset) output = strategy.unwrap( strategy.run(_custom_step, args=(next(iterator)))) return output dataset = strategy.experimental_distribute_dataset( distribute_strategy_test.get_dataset(strategy)) custom_function(dataset) writer.close() event_files = tf.io.gfile.glob( os.path.join(self.summary_dir, 'event*')) events_count_dictionary = { ('logits'): 0, } self.validate_recorded_sumary_file(event_files, events_count_dictionary, 1)