def test_cpu(self):
"""Run full training for MNIST CPU training."""
model_dir = self.get_tmp_model_dir()
start_time = time.time()
mnist_lib.train_and_evaluate(
config=config_lib.get_config(), model_dir=model_dir)
benchmark_time = time.time() - start_time
summaries = self.read_summaries(model_dir)
# Summaries contain all the information necessary for the regression
# metrics.
wall_time, _, eval_accuracy = zip(*summaries['eval_accuracy'])
wall_time = np.array(wall_time)
sec_per_epoch = np.mean(wall_time[1:] - wall_time[:-1])
end_eval_accuracy = eval_accuracy[-1]
# Assertions are deferred until the test finishes, so the metrics are
# always reported and benchmark success is determined based on *all*
# assertions.
self.assertBetween(end_eval_accuracy, 0.98, 1.0)
# Use the reporting API to report single or multiple metrics/extras.
self.report_wall_time(benchmark_time)
self.report_metrics({
'sec_per_epoch': sec_per_epoch,
'accuracy': end_eval_accuracy,
})
self.report_extras({
'model_name': 'MNIST',
'description': 'CPU test for MNIST.'
})
def main(_):
# Make sure tf does not allocate gpu memory.
tf.config.experimental.set_visible_devices([], 'GPU')
mnist_lib.train_and_evaluate(model_dir=FLAGS.model_dir,
num_epochs=FLAGS.num_epochs,
batch_size=FLAGS.batch_size,
learning_rate=FLAGS.learning_rate,
momentum=FLAGS.momentum)
def test_train_and_evaluate(self):
"""Runs a single train/eval step with mocked data."""
# Create a temporary directory where tensorboard metrics are written.
model_dir = tempfile.mkdtemp()
# Go two directories up to the root of the flax directory.
flax_root_dir = pathlib.Path(__file__).parents[2]
data_dir = str(flax_root_dir) + '/.tfds/metadata'
with tfds.testing.mock_data(num_examples=8, data_dir=data_dir):
mnist_lib.train_and_evaluate(model_dir=model_dir,
num_epochs=1,
batch_size=8,
learning_rate=0.1,
momentum=0.9)
def test_train_and_evaluate(self):
"""Runs a single train/eval step with mocked data."""
# Create a temporary directory where tensorboard metrics are written.
model_dir = tempfile.mkdtemp()
# Go two directories up to the root of the flax directory.
flax_root_dir = pathlib.Path(__file__).parents[2]
data_dir = str(flax_root_dir) + '/.tfds/metadata'
# Define training configuration.
config = config_lib.get_config()
config.num_epochs = 1
config.batch_size = 8
with tfds.testing.mock_data(num_examples=8, data_dir=data_dir):
mnist_lib.train_and_evaluate(config=config, model_dir=model_dir)
def main(_):
mnist_lib.train_and_evaluate(model_dir=FLAGS.model_dir,
num_epochs=FLAGS.num_epochs,
batch_size=FLAGS.batch_size,
learning_rate=FLAGS.learning_rate,
momentum=FLAGS.momentum)
def main(_):
# Make sure tf does not allocate gpu memory.
tf.config.experimental.set_visible_devices([], 'GPU')
mnist_lib.train_and_evaluate(config=FLAGS.config,
model_dir=FLAGS.model_dir)
