def main(unused_argv): # Necessary to use the tfds loader. tf.enable_v2_behavior() if jax.process_count() > 1: # TODO(ankugarg): Add support for multihost inference. raise NotImplementedError( 'BLEU eval does not support multihost inference.') rng = jax.random.PRNGKey(FLAGS.seed) mt_eval_config = json.loads(FLAGS.mt_eval_config) if FLAGS.experiment_config_filename: with tf.io.gfile.GFile(FLAGS.experiment_config_filename) as f: experiment_config = json.load(f) if jax.process_index() == 0: logging.info('experiment_config: %r', experiment_config) dataset_name = experiment_config['dataset'] model_name = experiment_config['model'] else: assert FLAGS.dataset and FLAGS.model dataset_name = FLAGS.dataset model_name = FLAGS.model if jax.process_index() == 0: logging.info('argv:\n%s', ' '.join(sys.argv)) logging.info('device_count: %d', jax.device_count()) logging.info('num_hosts : %d', jax.host_count()) logging.info('host_id : %d', jax.host_id()) model_class = models.get_model(model_name) dataset_builder = datasets.get_dataset(dataset_name) dataset_meta_data = datasets.get_dataset_meta_data(dataset_name) hparam_overrides = None if FLAGS.hparam_overrides: if isinstance(FLAGS.hparam_overrides, str): hparam_overrides = json.loads(FLAGS.hparam_overrides) merged_hps = hyperparameters.build_hparams( model_name=model_name, initializer_name=experiment_config['initializer'], dataset_name=dataset_name, hparam_file=FLAGS.trial_hparams_filename, hparam_overrides=hparam_overrides) if jax.process_index() == 0: logging.info('Merged hps are: %s', json.dumps(merged_hps.to_json())) evaluator = bleu_evaluator.BLEUEvaluator(FLAGS.checkpoint_dir, merged_hps, rng, model_class, dataset_builder, dataset_meta_data, mt_eval_config) evaluator.translate_and_calculate_bleu()
def main(unused_argv): # Necessary to use the tfds imagenet loader. tf.enable_v2_behavior() rng = jax.random.PRNGKey(FLAGS.seed) if FLAGS.hessian_eval_config: hessian_eval_config = json.loads(FLAGS.hessian_eval_config) else: hessian_eval_config = hessian_eval.DEFAULT_EVAL_CONFIG if FLAGS.experiment_config_filename: with tf.io.gfile.GFile(FLAGS.experiment_config_filename, 'r') as f: experiment_config = json.load(f) if jax.process_index() == 0: logging.info('experiment_config: %r', experiment_config) dataset_name = experiment_config['dataset'] model_name = experiment_config['model'] else: assert FLAGS.dataset and FLAGS.model dataset_name = FLAGS.dataset model_name = FLAGS.model if jax.process_index() == 0: logging.info('argv:\n%s', ' '.join(sys.argv)) logging.info('device_count: %d', jax.device_count()) logging.info('num_hosts : %d', jax.process_count()) logging.info('host_id : %d', jax.process_index()) model = models.get_model(model_name) dataset_builder = datasets.get_dataset(dataset_name) dataset_meta_data = datasets.get_dataset_meta_data(dataset_name) with tf.io.gfile.GFile(FLAGS.trial_hparams_filename, 'r') as f: hps = config_dict.ConfigDict(json.load(f)) if FLAGS.hparam_overrides: if isinstance(FLAGS.hparam_overrides, str): hparam_overrides = json.loads(FLAGS.hparam_overrides) hps.update_from_flattened_dict(hparam_overrides) run_lanczos.eval_checkpoints( FLAGS.checkpoint_dir, hps, rng, FLAGS.eval_num_batches, model, dataset_builder, dataset_meta_data, hessian_eval_config, FLAGS.min_global_step, FLAGS.max_global_step)
def _run(train_fn, dataset_name, eval_batch_size, eval_num_batches, eval_train_num_batches, eval_frequency, checkpoint_steps, num_tf_data_prefetches, num_device_prefetches, num_tf_data_map_parallel_calls, early_stopping_target_name, early_stopping_target_value, early_stopping_mode, eval_steps, hparam_file, hparam_overrides, initializer_name, model_name, loss_name, metrics_name, num_train_steps, experiment_dir, worker_id, training_metrics_config, callback_configs, external_checkpoint_path): """Function that runs a Jax experiment. See flag definitions for args.""" model_cls = models.get_model(model_name) initializer = initializers.get_initializer(initializer_name) dataset_builder = datasets.get_dataset(dataset_name) dataset_meta_data = datasets.get_dataset_meta_data(dataset_name) input_pipeline_hps = config_dict.ConfigDict( dict( num_tf_data_prefetches=num_tf_data_prefetches, num_device_prefetches=num_device_prefetches, num_tf_data_map_parallel_calls=num_tf_data_map_parallel_calls, )) merged_hps = hyperparameters.build_hparams( model_name=model_name, initializer_name=initializer_name, dataset_name=dataset_name, hparam_file=hparam_file, hparam_overrides=hparam_overrides, input_pipeline_hps=input_pipeline_hps) # Note that one should never tune an RNG seed!!! The seed is only included in # the hparams for convenience of running hparam trials with multiple seeds per # point. rng_seed = merged_hps.rng_seed if merged_hps.rng_seed < 0: rng_seed = _create_synchronized_rng_seed() xm_experiment = None xm_work_unit = None if jax.process_index() == 0: logging.info('Running with seed %d', rng_seed) rng = jax.random.PRNGKey(rng_seed) # Build the loss_fn, metrics_bundle, and flax_module. model = model_cls(merged_hps, dataset_meta_data, loss_name, metrics_name) trial_dir = os.path.join(experiment_dir, str(worker_id)) meta_data_path = os.path.join(trial_dir, 'meta_data.json') meta_data = {'worker_id': worker_id, 'status': 'incomplete'} if jax.process_index() == 0: logging.info('rng: %s', rng) gfile.makedirs(trial_dir) # Set up the metric loggers for host 0. metrics_logger, init_logger = utils.set_up_loggers( trial_dir, xm_work_unit) hparams_fname = os.path.join(trial_dir, 'hparams.json') logging.info('saving hparams to %s', hparams_fname) with gfile.GFile(hparams_fname, 'w') as f: f.write(merged_hps.to_json()) _write_trial_meta_data(meta_data_path, meta_data) else: metrics_logger = None init_logger = None try: epoch_reports = list( train_fn(trial_dir, model, dataset_builder, initializer, num_train_steps, merged_hps, rng, eval_batch_size, eval_num_batches, eval_train_num_batches, eval_frequency, checkpoint_steps, early_stopping_target_name, early_stopping_target_value, early_stopping_mode, eval_steps, metrics_logger, init_logger, training_metrics_config=training_metrics_config, callback_configs=callback_configs, external_checkpoint_path=external_checkpoint_path)) logging.info(epoch_reports) meta_data['status'] = 'done' except utils.TrainingDivergedError as err: meta_data['status'] = 'diverged' raise err finally: if jax.process_index() == 0: _write_trial_meta_data(meta_data_path, meta_data)
def test_early_stopping(self): """Test training early stopping on MNIST with a small model.""" rng = jax.random.PRNGKey(0) # Set the numpy seed to make the fake data deterministc. mocking.mock_data # ultimately calls numpy.random. np.random.seed(0) model_name = 'fully_connected' loss_name = 'cross_entropy' metrics_name = 'classification_metrics' initializer_name = 'noop' dataset_name = 'mnist' model_cls = models.get_model(model_name) initializer = initializers.get_initializer(initializer_name) dataset_builder = datasets.get_dataset(dataset_name) hparam_overrides = { 'lr_hparams': { 'base_lr': 0.1, 'schedule': 'cosine' }, 'batch_size': 8, 'train_size': 160, 'valid_size': 96, 'test_size': 80, } input_pipeline_hps = config_dict.ConfigDict( dict( num_tf_data_prefetches=-1, num_device_prefetches=0, num_tf_data_map_parallel_calls=-1, )) hps = hyperparameters.build_hparams( model_name, initializer_name, dataset_name, hparam_file=None, hparam_overrides=hparam_overrides, input_pipeline_hps=input_pipeline_hps) eval_batch_size = 16 num_examples = 256 def as_dataset(self, *args, **kwargs): del args del kwargs # pylint: disable=g-long-lambda,g-complex-comprehension return tf.data.Dataset.from_generator( lambda: ({ 'image': np.ones(shape=(28, 28, 1), dtype=np.uint8), 'label': 9, } for i in range(num_examples)), output_types=self.info.features.dtype, output_shapes=self.info.features.shape, ) # This will override the tfds.load(mnist) call to return 100 fake samples. with tfds.testing.mock_data(as_dataset_fn=as_dataset, num_examples=num_examples): dataset = dataset_builder(shuffle_rng=jax.random.PRNGKey(0), batch_size=hps.batch_size, eval_batch_size=eval_batch_size, hps=hps) model = model_cls(hps, datasets.get_dataset_meta_data(dataset_name), loss_name, metrics_name) num_train_steps = 40 early_stopping_target_name = 'test/ce_loss' early_stopping_target_value = 0.005 early_stopping_mode = 'less' eval_num_batches = 5 eval_every = 10 checkpoint_steps = [1, 3, 15] metrics_logger, init_logger = utils.set_up_loggers(self.test_dir) epoch_reports = list( trainer.train( train_dir=self.test_dir, model=model, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, initializer=initializer, num_train_steps=num_train_steps, hps=hps, rng=rng, eval_batch_size=eval_batch_size, eval_num_batches=eval_num_batches, eval_train_num_batches=eval_num_batches, eval_frequency=eval_every, checkpoint_steps=checkpoint_steps, early_stopping_target_name=early_stopping_target_name, early_stopping_target_value=early_stopping_target_value, early_stopping_mode=early_stopping_mode, metrics_logger=metrics_logger, init_logger=init_logger)) self.assertLen(epoch_reports, 3) self.assertGreater(epoch_reports[-2][early_stopping_target_name], early_stopping_target_value) self.assertLess(epoch_reports[-1][early_stopping_target_name], early_stopping_target_value)
def test_trainer(self): """Test training for two epochs on MNIST with a small model.""" rng = jax.random.PRNGKey(0) # Set the numpy seed to make the fake data deterministc. mocking.mock_data # ultimately calls numpy.random. np.random.seed(0) model_name = 'fully_connected' loss_name = 'cross_entropy' metrics_name = 'classification_metrics' initializer_name = 'noop' dataset_name = 'mnist' model_cls = models.get_model(model_name) initializer = initializers.get_initializer(initializer_name) dataset_builder = datasets.get_dataset(dataset_name) hparam_overrides = { 'lr_hparams': { 'base_lr': 0.1, 'schedule': 'cosine' }, 'batch_size': 8, 'train_size': 160, 'valid_size': 96, 'test_size': 80, } input_pipeline_hps = config_dict.ConfigDict( dict( num_tf_data_prefetches=-1, num_device_prefetches=0, num_tf_data_map_parallel_calls=-1, )) hps = hyperparameters.build_hparams( model_name, initializer_name, dataset_name, hparam_file=None, hparam_overrides=hparam_overrides, input_pipeline_hps=input_pipeline_hps) eval_batch_size = 16 num_examples = 256 def as_dataset(self, *args, **kwargs): del args del kwargs # pylint: disable=g-long-lambda,g-complex-comprehension return tf.data.Dataset.from_generator( lambda: ({ 'image': np.ones(shape=(28, 28, 1), dtype=np.uint8), 'label': 9, } for i in range(num_examples)), output_types=self.info.features.dtype, output_shapes=self.info.features.shape, ) # This will override the tfds.load(mnist) call to return 100 fake samples. with tfds.testing.mock_data(as_dataset_fn=as_dataset, num_examples=num_examples): dataset = dataset_builder(shuffle_rng=jax.random.PRNGKey(0), batch_size=hps.batch_size, eval_batch_size=eval_batch_size, hps=hps) model = model_cls(hps, datasets.get_dataset_meta_data(dataset_name), loss_name, metrics_name) num_train_steps = 40 eval_num_batches = 5 eval_every = 10 checkpoint_steps = [1, 3, 15] metrics_logger, init_logger = utils.set_up_loggers(self.test_dir) epoch_reports = list( trainer.train( train_dir=self.test_dir, model=model, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, initializer=initializer, num_train_steps=num_train_steps, hps=hps, rng=rng, eval_batch_size=eval_batch_size, eval_num_batches=eval_num_batches, eval_train_num_batches=eval_num_batches, eval_frequency=eval_every, checkpoint_steps=checkpoint_steps, metrics_logger=metrics_logger, init_logger=init_logger)) # check that the additional checkpoints are saved. checkpoint_dir = os.path.join(self.test_dir, 'checkpoints') saved_steps = [] for f in tf.io.gfile.listdir(checkpoint_dir): if f[:5] == 'ckpt_': saved_steps.append(int(f[5:])) self.assertEqual(set(saved_steps), set(checkpoint_steps)) self.assertLen(epoch_reports, num_train_steps / eval_every) with tf.io.gfile.GFile(os.path.join(self.test_dir, 'measurements.csv')) as f: df = pandas.read_csv(f) train_err = df['train/error_rate'].values[-1] self.assertEqual(df['preemption_count'].values[-1], 0) self.assertLess(train_err, 0.9) self.assertEqual(set(df.columns.values), set(get_column_names())) model = model_cls(hps, {'apply_one_hot_in_loss': False}, loss_name, metrics_name) # Test reload from the checkpoint by increasing num_train_steps. num_train_steps_reload = 100 epoch_reports = list( trainer.train( train_dir=self.test_dir, model=model, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, initializer=initializer, num_train_steps=num_train_steps_reload, hps=hps, rng=rng, eval_batch_size=eval_batch_size, eval_num_batches=eval_num_batches, eval_train_num_batches=eval_num_batches, eval_frequency=eval_every, checkpoint_steps=checkpoint_steps, metrics_logger=metrics_logger, init_logger=init_logger)) self.assertLen(epoch_reports, (num_train_steps_reload - num_train_steps) / eval_every) with tf.io.gfile.GFile(os.path.join(self.test_dir, 'measurements.csv')) as f: df = pandas.read_csv(f) train_err = df['train/error_rate'].values[-1] train_loss = df['train/ce_loss'].values[-1] self.assertLess(train_err, 0.35) self.assertLess(train_loss, 0.1) self.assertEqual(df['valid/num_examples'].values[-1], eval_num_batches * eval_batch_size) self.assertEqual(df['preemption_count'].values[-1], 1) # Check that the correct learning rate was saved in the measurements file. final_learning_rate = df['learning_rate'].values[-1] final_step = df['global_step'].values[-1] self.assertEqual(num_train_steps_reload, final_step) # final_step will be one larger than the last step used to calculate the # lr_decay, hense we plug in (final_step - 1) to the decay formula. # Note that there is a small numerical different here with np vs jnp. decay_factor = (1 + np.cos( (final_step - 1) / num_train_steps_reload * np.pi)) * 0.5 self.assertEqual(float(final_learning_rate), hps.lr_hparams['base_lr'] * decay_factor) self.assertEqual(set(df.columns.values), set(get_column_names()))
def test_run_lanczos(self): """Test training for two epochs on MNIST with a small model.""" rng = jax.random.PRNGKey(0) # Set the numpy seed to make the fake data deterministc. mocking.mock_data # ultimately calls numpy.random. np.random.seed(0) model_name = 'fully_connected' loss_name = 'cross_entropy' metrics_name = 'classification_metrics' initializer_name = 'noop' dataset_name = 'mnist' model_cls = models.get_model(model_name) initializer = initializers.get_initializer(initializer_name) dataset_builder = datasets.get_dataset(dataset_name) hparam_overrides = { 'lr_hparams': { 'base_lr': 0.1, 'schedule': 'cosine' }, 'batch_size': 8, 'train_size': 160, 'valid_size': 96, 'test_size': 80, } input_pipeline_hps = config_dict.ConfigDict(dict( num_tf_data_prefetches=-1, num_device_prefetches=0, num_tf_data_map_parallel_calls=-1, )) hps = hyperparameters.build_hparams( model_name, initializer_name, dataset_name, hparam_file=None, hparam_overrides=hparam_overrides, input_pipeline_hps=input_pipeline_hps) model = model_cls(hps, datasets.get_dataset_meta_data(dataset_name), loss_name, metrics_name) eval_batch_size = 16 num_examples = 256 def as_dataset(self, *args, **kwargs): del args del kwargs # pylint: disable=g-long-lambda,g-complex-comprehension return tf.data.Dataset.from_generator( lambda: ({ 'image': np.ones(shape=(28, 28, 1), dtype=np.uint8), 'label': 9, } for i in range(num_examples)), output_types=self.info.features.dtype, output_shapes=self.info.features.shape, ) # This will override the tfds.load(mnist) call to return 100 fake samples. with tfds.testing.mock_data( as_dataset_fn=as_dataset, num_examples=num_examples): dataset = dataset_builder( shuffle_rng=jax.random.PRNGKey(0), batch_size=hps.batch_size, eval_batch_size=eval_batch_size, hps=hps) num_train_steps = 41 eval_num_batches = 5 eval_every = 10 checkpoint_steps = [10, 30, 40] metrics_logger, init_logger = None, None _ = list( trainer.train( train_dir=self.test_dir, model=model, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, initializer=initializer, num_train_steps=num_train_steps, hps=hps, rng=rng, eval_batch_size=eval_batch_size, eval_num_batches=eval_num_batches, eval_train_num_batches=eval_num_batches, eval_frequency=eval_every, checkpoint_steps=checkpoint_steps, metrics_logger=metrics_logger, init_logger=init_logger)) checkpoint_dir = os.path.join(self.test_dir, 'checkpoints') rng = jax.random.PRNGKey(0) run_lanczos.eval_checkpoints( checkpoint_dir, hps, rng, eval_num_batches, model_cls=model_cls, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, dataset_meta_data=datasets.get_dataset_meta_data(dataset_name), hessian_eval_config=hessian_eval.DEFAULT_EVAL_CONFIG, ) # Load the saved file. hessian_dir = os.path.join(checkpoint_dir, 'hessian', 'training_metrics') pytree_list = checkpoint.load_pytree(hessian_dir) # Convert to a regular list (checkpointer will have converted the saved # list to a dict of keys '0', '1', ... pytree_list = [pytree_list[str(i)] for i in range(len(pytree_list))] # Test that the logged steps are correct. saved_steps = [row['step'] for row in pytree_list] self.assertEqual(saved_steps, checkpoint_steps)
def test_shampoo_wrn(self): """Test distributed shampoo on fake dataset.""" model_name = 'simple_cnn' model_cls = models.get_model(model_name) hparam_overrides = { 'optimizer': 'distributed_shampoo', 'batch_size': 1, 'train_size': 10, 'valid_size': 10, 'input_shape': (32, 32, 3), 'output_shape': (10,), 'opt_hparams': { 'block_size': 32, 'beta1': 0.9, 'beta2': 0.999, 'diagonal_epsilon': 1e-10, 'matrix_epsilon': 1e-6, 'weight_decay': 0.0, 'start_preconditioning_step': 5, 'preconditioning_compute_steps': 1, 'statistics_compute_steps': 1, 'best_effort_shape_interpretation': True, 'graft_type': distributed_shampoo.GraftingType.SGD, 'nesterov': True, 'exponent_override': 0, 'batch_axis_name': 'batch', 'num_devices_for_pjit': None, 'shard_optimizer_states': False, 'inverse_failure_threshold': 0.1, 'clip_by_scaled_gradient_norm': None, 'precision': lax.Precision.HIGHEST, 'moving_average_for_momentum': False, 'skip_preconditioning_dim_size_gt': 4096, 'best_effort_memory_usage_reduction': False, }, } input_pipeline_hps = config_dict.ConfigDict(dict( num_tf_data_prefetches=-1, num_device_prefetches=0, num_tf_data_map_parallel_calls=-1, )) hps = hyperparameters.build_hparams( model_name, initializer_name='noop', dataset_name='fake', hparam_file=None, hparam_overrides=hparam_overrides, input_pipeline_hps=input_pipeline_hps) initializer = initializers.get_initializer('noop') dataset_builder = datasets.get_dataset('fake') dataset = dataset_builder( shuffle_rng=jax.random.PRNGKey(0), batch_size=hps.batch_size, eval_batch_size=hps.batch_size, hps=hps) loss_name = 'cross_entropy' metrics_name = 'classification_metrics' dataset_meta_data = datasets.get_dataset_meta_data('fake') model = model_cls(hps, dataset_meta_data, loss_name, metrics_name) metrics_logger, init_logger = utils.set_up_loggers(self.test_dir) _ = list( trainer.train( train_dir=self.test_dir, model=model, dataset_builder=lambda *unused_args, **unused_kwargs: dataset, initializer=initializer, num_train_steps=1, hps=hps, rng=jax.random.PRNGKey(42), eval_batch_size=hps.batch_size, eval_num_batches=None, eval_train_num_batches=None, eval_frequency=10, checkpoint_steps=[], metrics_logger=metrics_logger, init_logger=init_logger)) with tf.io.gfile.GFile(os.path.join(self.test_dir, 'measurements.csv')) as f: df = pandas.read_csv(f) valid_ce_loss = df['valid/ce_loss'].values[-1] self.assertLess(valid_ce_loss, 1e-3)