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 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_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_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)