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)