def delete_backup(self):
"""Delete the backup directories.
Delete the backup directories which should not exist after `fit()`
successfully finishes.
"""
# pylint: disable=protected-access
for pathname in file_io.get_matching_files_v2(
self.write_checkpoint_manager._prefix + '*'):
file_io.delete_recursively_v2(pathname)
for pathname in file_io.get_matching_files_v2(
os.path.join(self.write_checkpoint_manager.directory, 'checkpoint')):
file_io.delete_recursively_v2(pathname)
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 = file_io.get_matching_files_v2(
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 testV2SummaryWithKerasSequentialModel(self):
strategy = get_tpu_strategy()
with strategy.scope():
model = mnist_model((28, 28, 3))
model.compile('sgd', 'mse')
dataset = get_image_dataset()
tensorboard_callback = callbacks.TensorBoard(self.summary_dir,
update_freq=2)
model.fit(dataset,
steps_per_epoch=10,
epochs=1,
callbacks=[tensorboard_callback])
events_count_dictionary = {
'sequential/layer_for_histogram_summary/custom_histogram_summary_v2':
0,
'sequential/layer_for_image_summary/custom_image_summary_v2':
0,
}
event_files = file_io.get_matching_files_v2(
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.
self.validate_recorded_sumary_file(event_files,
events_count_dictionary, 5)
def testV2SummaryWithKerasFit(self):
resolver = tpu_cluster_resolver.TPUClusterResolver('')
tpu_strategy_util.initialize_tpu_system(resolver)
strategy = tpu_strategy_lib.TPUStrategy(resolver)
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 = file_io.get_matching_files_v2(
os.path.join(self.summary_dir, 'train', 'event*'))
events_count_dictionary = {
'custom_model/layer_for_scalar_summary/custom_scalar_summary':
0,
'custom_model/layer_for_histogram_summary/custom_histogram_summary':
0
}
for event_file in event_files:
for e in summary_iterator.summary_iterator(event_file):
for v in e.summary.value:
if v.tag in events_count_dictionary:
events_count_dictionary[v.tag] += 1
# 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.assertEqual(
events_count_dictionary[(
'custom_model/layer_for_histogram_summary/'
'custom_histogram_summary')], 5)
self.assertEqual(
events_count_dictionary[
'custom_model/layer_for_scalar_summary/custom_scalar_summary'],
5)
def testV2SummaryWithImageModel(self):
strategy = get_tpu_strategy()
with strategy.scope():
model = mnist_model((28, 28, 3))
model.compile('sgd', 'mse')
dataset = get_image_dataset()
tensorboard_callback = callbacks.TensorBoard(self.summary_dir,
update_freq=2)
model.fit(dataset,
steps_per_epoch=10,
epochs=1,
callbacks=[tensorboard_callback])
event_files = file_io.get_matching_files_v2(
os.path.join(self.summary_dir, 'train', 'event*'))
events_count_dictionary = {
('sequential/layer_for_histogram_summary'
'/custom_histogram_summary'):
0,
'sequential/layer_for_image_summary/custom_image_summary/image/0':
0
}
for event_file in event_files:
for e in summary_iterator.summary_iterator(event_file):
for v in e.summary.value:
if v.tag in events_count_dictionary:
events_count_dictionary[v.tag] += 1
# 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.assertEqual(
events_count_dictionary[(
'sequential/layer_for_histogram_summary/'
'custom_histogram_summary')], 5)
self.assertEqual(
events_count_dictionary[('sequential/layer_for_image_summary/'
'custom_image_summary/image/0')], 5)
def testSummaryWithCustomTrainingLoop(self):
strategy = get_tpu_strategy()
writer = summary_ops_v2.create_file_writer_v2(self.summary_dir)
with strategy.scope():
model = distribute_strategy_test.get_model()
model.compile('sgd', 'mse')
@def_function.function
def custom_function(dataset):
def _custom_step(features, labels):
del labels
logits = model(features)
with summary_ops_v2.record_if(True), writer.as_default():
scalar_summary_v2.scalar(
'logits',
math_ops.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 = file_io.get_matching_files_v2(
os.path.join(self.summary_dir, 'event*'))
events_count_dictionary = {
('logits'): 0,
}
self.validate_recorded_sumary_file(event_files, events_count_dictionary,
1)
