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)
