def test_schedule_stretching(self): """Test that schedules can be properly stretched.""" max_training_steps = 100 lr_hparams = config_dict.ConfigDict({ 'schedule': 'mlperf_polynomial', 'base_lr': 10.0, 'warmup_steps': 10, 'decay_end': -1, 'end_lr': 1e-4, 'power': 2.0, 'start_lr': 0.0, 'warmup_power': 1.0, }) lr_fn = schedules.get_schedule_fn(lr_hparams, max_training_steps) stretch_factor = 3 stretched_lr_fn = schedules.get_schedule_fn( lr_hparams, max_training_steps, stretch_factor=stretch_factor) lrs = [lr_fn(t) for t in range(max_training_steps)] stretched_lrs = [ stretched_lr_fn(t) for t in range(stretch_factor * max_training_steps) ] self.assertEqual(lrs, stretched_lrs[::stretch_factor]) self.assertEqual(lrs, stretched_lrs[1::stretch_factor]) self.assertEqual(lrs, stretched_lrs[2::stretch_factor]) # Assert that the stretched schedule has proper staircase behavior. for update_step in range(max_training_steps): start = update_step * stretch_factor end = (update_step + 1) * stretch_factor expected = [lrs[update_step]] * stretch_factor self.assertEqual(stretched_lrs[start:end], expected)
def test_raises(self): """Test that an exception is raised with extra hparams.""" good_hps = config_dict.ConfigDict( dict( lr_hparams={ 'schedule': 'mlperf_polynomial', 'base_lr': .1, 'warmup_steps': 200, 'decay_end': -1, 'end_lr': 1e-4, 'power': 2.0, 'start_lr': 0.0, 'warmup_power': 1.0, })) bad_hps = config_dict.ConfigDict( dict( lr_hparams={ 'schedule': 'mlperf_polynomial', 'base_lr': .1, 'warmup_steps': 200, 'initial_value': .1, 'decay_end': -1, 'end_lr': 1e-4, 'power': 2.0, 'start_lr': 0.0, })) bad_hps2 = config_dict.ConfigDict( dict( lr_hparams={ 'schedule': 'polynomial', 'power': 2.0, 'initial_value': .1, 'end_factor': .01, 'decay_steps': 200, 'decay_steps_factor': 0.5 })) # This should pass. schedules.get_schedule_fn(good_hps.lr_hparams, 1) # This should raise an exception due to the extra hparam. with self.assertRaises(ValueError): schedules.get_schedule_fn(bad_hps.lr_hparams, 1) # This should raise an exception due to the mutually exclusive hparams. with self.assertRaises(ValueError): schedules.get_schedule_fn(bad_hps2.lr_hparams, 1)
def test_polynomial_decay_decay_steps(self): """Test polynomial schedule works correctly with decay_steps.""" hps = config_dict.ConfigDict( dict( lr_hparams={ 'schedule': 'polynomial', 'power': 2.0, 'initial_value': .1, 'end_factor': .01, 'decay_steps': 200, })) max_training_steps = 400 lr_fn = schedules.get_schedule_fn(hps.lr_hparams, max_training_steps) hps = hps.lr_hparams decay_steps = hps['decay_steps'] for step in range(max_training_steps): expected_learning_rate = tf.train.polynomial_decay( hps['initial_value'], step, decay_steps, hps['end_factor'] * hps['initial_value'], power=hps['power'])().numpy() self.assertAlmostEqual(lr_fn(step), expected_learning_rate)
def train(train_dir, model, dataset_builder, initializer, num_train_steps, hps, rng, eval_batch_size, eval_num_batches, eval_train_num_batches, eval_frequency, checkpoint_steps, early_stopping_target_name=None, early_stopping_target_value=None, early_stopping_mode=None, eval_steps=None, metrics_logger=None, init_logger=None, training_metrics_config=None, callback_configs=None, external_checkpoint_path=None): """Main training loop. Trains the given network on the specified dataset for the given number of epochs. Saves the training curve in train_dir/r=3/results.tsv. Args: train_dir: (str) Path of the training directory. model: (BaseModel) Model object to be trained. dataset_builder: dataset builder returned by datasets.get_dataset. initializer: Must have API as defined in initializers.py num_train_steps: (int) Number of steps to train on. hps: (tf.HParams) Model, initialization and training hparams. rng: (jax.random.PRNGKey) Rng seed used in model initialization and data shuffling. eval_batch_size: the evaluation batch size. If None, use hps.batch_size. eval_num_batches: (int) The number of batches used for evaluating on validation and test sets. Set to None to evaluate on the whole train set. eval_train_num_batches: (int) The number of batches for evaluating on train. Set to None to evaluate on the whole training set. eval_frequency: (int) Evaluate every k steps. checkpoint_steps: List of integers indicating special steps to save checkpoints at. These checkpoints do not get used for preemption recovery. early_stopping_target_name: A string naming the metric to use to perform early stopping. If this metric reaches the value `early_stopping_target_value`, training will stop. Must include the dataset split (ex: validation/error_rate). early_stopping_target_value: A float indicating the value at which to stop training. early_stopping_mode: One of "above" or "below", indicates if we should stop when the metric is above or below the threshold value. Example: if "above", then training will stop when `report[early_stopping_target_name] >= early_stopping_target_value`. eval_steps: List of integers indicating which steps to perform evals. If provided, eval_frequency will be ignored. Performing an eval implies saving a checkpoint that will be used to resume training in the case of preemption. metrics_logger: Used to log all eval metrics during training. See utils.MetricLogger for API definition. init_logger: Used for black box initializers that have learning curves. training_metrics_config: Dict specifying the configuration of the training_metrics_grabber. Set to None to skip logging of advanced training metrics. callback_configs: List of configs specifying general callbacks to run during the eval phase. Empty list means no callbacks are run. See callbacks.py for details on what is expected in a config. external_checkpoint_path: (str) If this argument is set, we will load the optimizer_state, params, batch_stats, and training_metrics from the checkpoint at this location. Yields: metrics: A dictionary of all eval metrics from the given epoch. """ # NOTE: the initialization RNG should *not* be per-host, as this will create # different sets of weights per host. However, all other RNGs should be # per-host. # TODO(znado,gilmer,gdahl): implement replicating the same initialization # across hosts. rng, init_rng = jax.random.split(rng) rng = jax.random.fold_in(rng, jax.process_index()) rng, data_rng = jax.random.split(rng) # only used if checkpoints_steps is non-empty. checkpoint_dir = os.path.join(train_dir, 'checkpoints') # For logging / processing off the main thread pool = multiprocessing.pool.ThreadPool() if jax.process_index() == 0: logging.info('Let the training begin!') logging.info('Dataset input shape: %r', hps.input_shape) logging.info('Hyperparameters: %s', hps) if eval_batch_size is None: eval_batch_size = hps.batch_size if callback_configs is None: callback_configs = [] # Maybe run the initializer. unreplicated_params, unreplicated_batch_stats = init_utils.initialize( model.flax_module, initializer, model.loss_fn, hps.input_shape, hps.output_shape, hps, init_rng, init_logger, model.get_fake_batch(hps)) if jax.process_index() == 0: utils.log_pytree_shape_and_statistics(unreplicated_params) logging.info('train_size: %d,', hps.train_size) # Note that global_step refers to the number of gradients calculations, not # the number of model updates. This means when using gradient accumulation, # one must supply configs where the number of steps are in units of gradient # calculations, not model updates, and in post processing one must divide # global_step by grad_accum_step_multiplier to get the number of updates. # # If using gradient accumulation, stretch the learning rate schedule by the # number of gradient calculations per weight update. stretch_factor = 1 if hps.get('total_accumulated_batch_size') is not None: stretch_factor = hps.total_accumulated_batch_size // hps.batch_size lr_fn = schedules.get_schedule_fn(hps.lr_hparams, num_train_steps, stretch_factor=stretch_factor) optimizer_init_fn, optimizer_update_fn = optimizers.get_optimizer( hps, model) unreplicated_optimizer_state = optimizer_init_fn(unreplicated_params) (unreplicated_metrics_state, metrics_update_fn, metrics_summary_fn) = None, None, None if training_metrics_config is not None: (metrics_init_fn, metrics_update_fn, metrics_summary_fn) = make_training_metrics(num_train_steps, **training_metrics_config) unreplicated_metrics_state = metrics_init_fn(unreplicated_params) (optimizer_state, params, batch_stats, metrics_state, global_step, sum_train_cost, preemption_count, is_restored) = checkpoint.replicate_and_maybe_restore_checkpoint( unreplicated_optimizer_state, unreplicated_params, unreplicated_batch_stats, unreplicated_metrics_state, train_dir=train_dir, external_checkpoint_path=external_checkpoint_path) if is_restored: preemption_count += 1 # Fold the restored step into the dataset RNG so that we will get a # different shuffle each time we restore, so that we do not repeat a # previous dataset ordering again after restoring. This is not the only # difference in shuffling each pre-emption, because we often times reshuffle # the input files each time in a non-deterministic manner. # # Note that if we are pre-empted more than once per epoch then we will # retrain more on the beginning of the training split, because each time we # restore we refill the shuffle buffer with the first `shuffle_buffer_size` # elements from the training split to continue training. # # Also note that for evaluating on the training split, because we are # reshuffling each time, we will get a new eval_train split each time we are # pre-empted. data_rng = jax.random.fold_in(data_rng, global_step) assert hps.batch_size % (jax.device_count()) == 0 assert eval_batch_size % (jax.device_count()) == 0 dataset = dataset_builder( data_rng, hps.batch_size, eval_batch_size=eval_batch_size, hps=hps, ) update_fn = functools.partial(update, training_cost=model.training_cost, grad_clip=hps.get('grad_clip'), optimizer_update_fn=optimizer_update_fn, metrics_update_fn=metrics_update_fn) # in_axes = ( # optimizer_state = 0, # params = 0, # batch_stats = 0, # metrics_state = 0, # batch = 0, # step = None, # lr = None, # rng = None, # local_device_index = 0, # running_train_cost = 0, # training_cost, # grad_clip, # optimizer_update_fn, # metrics_state_update_fn) # Also, we can donate buffers for 'optimizer', 'batch_stats', # 'batch' and 'training_metrics_state' for update's pmapped computation. update_pmapped = jax.pmap(update_fn, axis_name='batch', in_axes=(0, 0, 0, 0, 0, None, None, None, 0, 0), donate_argnums=(0, 1, 2, 8)) # During eval, we can donate the 'batch' buffer. We don't donate the # 'params' and 'batch_stats' buffers as we don't re-assign those values in # eval, we do that only in train. evaluate_batch_pmapped = jax.pmap(model.evaluate_batch, axis_name='batch', donate_argnums=(2, )) start_time = time.time() start_step = global_step prev_eval_step = start_step def get_step_frequency(cur_step): return float(cur_step - start_step) / (time.time() - start_time) if jax.process_index() == 0: trainer_utils.log_message('Starting training!', pool, xm_work_unit) # Numpy array of range(0, local_device_count) to send to each device to be # folded into the RNG inside each train step to get a unique per-device RNG. local_device_indices = np.arange(jax.local_device_count()) # Start at the resumed step and continue until we have finished the number of # training steps. If building a dataset iterator using a tf.data.Dataset, in # the case of a batch size that does not evenly divide the training dataset # size, if using `ds.batch(..., drop_remainer=True)` on the training dataset # then the final batch in this iterator will be a partial batch. However, if # `drop_remainer=False`, then this iterator will always return batches of the # same size, and the final batch will have elements from the start of the # (num_epochs + 1)-th epoch. train_iter = itertools.islice(dataset.train_iterator_fn(), global_step, num_train_steps) eval_callbacks = [] rng, callback_rng = jax.random.split(rng) callback_rngs = jax.random.split(callback_rng, len(callback_configs)) for callback_rng, config in zip(callback_rngs, callback_configs): eval_callback = callbacks.get_callback(config['callback_name'])( model, params, batch_stats, optimizer_state, dataset, hps, config, train_dir, callback_rng) eval_callbacks.append(eval_callback) train_iter = trainer_utils.prefetch_input_pipeline( train_iter, hps.num_device_prefetches) eval_start_time = start_time eval_start_step = start_step for _ in range(start_step, num_train_steps): with jax.profiler.StepTraceAnnotation('train', step_num=global_step): # NOTE(dsuo): to properly profile each step, we must include batch # creation in the StepTraceContext (as opposed to putting `train_iter` # directly in the top-level for loop). batch = next(train_iter) if global_step in checkpoint_steps and jax.process_index() == 0: checkpoint.save_unreplicated_checkpoint_background( checkpoint_dir, optimizer_state, params, batch_stats, metrics_state, global_step, preemption_count, sum_train_cost, max_to_keep=None) lr = lr_fn(global_step) optimizer_state, params, batch_stats, sum_train_cost, metrics_state, grad_norm = update_pmapped( optimizer_state, params, batch_stats, metrics_state, batch, global_step, lr, rng, local_device_indices, sum_train_cost) global_step += 1 # TODO(gdahl, gilmer): consider moving this test up. # NB: Since this test is after we increment global_step, having 0 in # eval_steps does nothing. if trainer_utils.should_eval(global_step, eval_frequency, eval_steps): train_steps_per_sec = (global_step - eval_start_step) / ( time.time() - eval_start_time) eval_start_step = global_step eval_start_time = time.time() batch_stats = trainer_utils.maybe_sync_batchnorm_stats( batch_stats) report, eval_time = eval_metrics(params, batch_stats, dataset, eval_num_batches, eval_train_num_batches, evaluate_batch_pmapped) mean_train_cost = sum_train_cost.mean().item() / max( 1, global_step - prev_eval_step) report.update( learning_rate=float(lr), global_step=global_step, epoch=global_step * hps.batch_size // hps.train_size, train_steps_per_sec=train_steps_per_sec, overall_steps_per_sec=get_step_frequency(global_step), eval_time=eval_time, grad_norm=np.mean(grad_norm), preemption_count=preemption_count, train_cost=mean_train_cost) for eval_callback in eval_callbacks: callback_metrics = eval_callback.run_eval( params, batch_stats, optimizer_state, global_step) if set(callback_metrics.keys()).intersection( set(report.keys())): raise ValueError( 'There was a collision between the callback' 'metrics and the standard eval metrics keys') report.update(callback_metrics) yield report if jax.process_index() == 0: trainer_utils.log_eta(pool, xm_work_unit, global_step, train_steps_per_sec, num_train_steps, start_time, eval_frequency, eval_steps, eval_time) trainer_utils.log_epoch_report(report, metrics_logger) trainer_utils.maybe_log_training_metrics( metrics_state, metrics_summary_fn, metrics_logger) checkpoint.save_unreplicated_checkpoint_background( train_dir, optimizer_state, params, batch_stats, metrics_state, global_step, preemption_count, sum_train_cost) sum_train_cost = jnp.zeros(jax.local_device_count()) prev_eval_step = global_step early_stopping_condition = trainer_utils.check_for_early_stopping( early_stopping_target_name, early_stopping_target_value, early_stopping_mode, report) if early_stopping_condition: comparison_string = '>=' if early_stopping_mode == 'above' else '<=' logging.info( 'Early stopping because metric %s=%f, reached the target value ' 'of %s %f.', early_stopping_target_name, report[early_stopping_target_name], comparison_string, early_stopping_target_value) return # Always log and checkpoint on host 0 at the end of training. # If we moved where in the loop body evals happen then we would not need this # test. if prev_eval_step != num_train_steps: train_steps_per_sec = (global_step - eval_start_step) / ( time.time() - eval_start_time) batch_stats = trainer_utils.maybe_sync_batchnorm_stats(batch_stats) report, eval_time = eval_metrics(params, batch_stats, dataset, eval_num_batches, eval_train_num_batches, evaluate_batch_pmapped) lr = lr_fn(global_step) # Correct the average for the final partial epoch. mean_train_cost = sum_train_cost.mean().item() / max( 1, global_step - prev_eval_step) report.update(learning_rate=float(lr), global_step=global_step, epoch=global_step * hps.batch_size // hps.train_size, train_steps_per_sec=train_steps_per_sec, overall_steps_per_sec=get_step_frequency(global_step), eval_time=eval_time, grad_norm=np.mean(grad_norm), preemption_count=preemption_count, train_cost=mean_train_cost) yield report if jax.process_index() == 0: trainer_utils.log_eta(pool, xm_work_unit, global_step, train_steps_per_sec, num_train_steps, start_time, eval_frequency, eval_steps, eval_time) trainer_utils.log_epoch_report(report, metrics_logger) trainer_utils.maybe_log_training_metrics(metrics_state, metrics_summary_fn, metrics_logger) checkpoint.save_unreplicated_checkpoint_background( train_dir, optimizer_state, params, batch_stats, metrics_state, global_step, preemption_count, sum_train_cost) # To make sure the last checkpoint was correctly saved. checkpoint.wait_for_checkpoint_save()
def test_mlperf_schedule(self): """Test there are no changes to the MLPerf polynomial decay schedule.""" expected_lrs = [ 0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.2, 1.4, 1.6, 1.8, 2.0, 2.2, 2.4, 2.6, 2.8, 3.0, 3.2, 3.4, 3.6, 3.8, 4.0, 4.2, 4.4, 4.6, 4.8, 5.0, 5.2, 5.4, 5.6, 5.8, 6.0, 6.2, 6.4, 6.6, 6.8, 7.0, 7.2, 7.4, 7.6, 7.8, 8.0, 8.2, 8.4, 8.6, 8.8, 9.0, 9.2, 9.4, 9.6, 9.8, 10.0, 9.802962, 9.607885, 9.414769, 9.223614, 9.034419, 8.847184, 8.661909, 8.478596, 8.297242, 8.117851, 7.940418, 7.764947, 7.5914364, 7.419886, 7.2502966, 7.082668, 6.917, 6.7532916, 6.591545, 6.4317584, 6.273932, 6.1180663, 5.964162, 5.812217, 5.662234, 5.5142093, 5.368148, 5.2240453, 5.0819044, 4.941723, 4.803503, 4.6672425, 4.532944, 4.4006047, 4.2702274, 4.1418095, 4.0153522, 3.8908558, 3.7683203, 3.647745, 3.5291305, 3.4124763, 3.297783, 3.18505, 3.0742776, 2.965466, 2.858614, 2.753724, 2.6507936, 2.5498245, 2.4508152, 2.353767, 2.2586792, 2.1655521, 2.0743854, 1.9851794, 1.8979341, 1.8126491, 1.7293249, 1.6479613, 1.568558, 1.4911155, 1.4156334, 1.342112, 1.2705511, 1.2009507, 1.133311, 1.0676318, 1.0039133, 0.94215524, 0.8823574, 0.8245205, 0.7686443, 0.71472853, 0.66277343, 0.61277884, 0.56474483, 0.5186714, 0.4745585, 0.43240622, 0.39221448, 0.35398334, 0.31771275, 0.28340274, 0.2510533, 0.22066444, 0.19223614, 0.16576843, 0.14126128, 0.11871469, 0.098128565, 0.079503134, 0.06283828, 0.048134, 0.035390284, 0.02460714, 0.01578457, 0.00892257, 0.004021142, ] hps = config_dict.ConfigDict( dict( lr_hparams={ 'schedule': 'mlperf_polynomial', 'base_lr': 10.0, 'warmup_steps': 50, 'decay_end': -1, 'end_lr': 1e-4, 'power': 2.0, 'start_lr': 0.0, 'warmup_power': 1.0, })) max_training_steps = 50 lr_fn = schedules.get_schedule_fn(hps.lr_hparams, max_training_steps) for step in range(max_training_steps): self.assertAlmostEqual(lr_fn(step), expected_lrs[step])
def train(train_dir, model, dataset_builder, initializer, num_train_steps, hps, rng, eval_batch_size, eval_num_batches, eval_train_num_batches, eval_frequency, checkpoint_steps, eval_steps=None, metrics_logger=None, init_logger=None, training_metrics_config=None, use_deprecated_checkpointing=True): """Main training loop. Trains the given network on the specified dataset for the given number of epochs. Saves the training curve in train_dir/r=3/results.tsv. Args: train_dir: (str) Path of the training directory. model: (BaseModel) Model object to be trained. dataset_builder: dataset builder returned by datasets.get_dataset. initializer: Must have API as defined in initializers.py num_train_steps: (int) Number of steps to train on. hps: (tf.HParams) Model, initialization and training hparams. rng: (jax.random.PRNGKey) Rng seed used in model initialization and data shuffling. eval_batch_size: the evaluation batch size. If None, use hps.batch_size. eval_num_batches: (int) The number of batches used for evaluating on validation and test sets. Set to None to evaluate on the whole train set. eval_train_num_batches: (int) The number of batches for evaluating on train. Set to None to evaluate on the whole training set. eval_frequency: (int) Evaluate every k steps. checkpoint_steps: List of integers indicating special steps to save checkpoints at. These checkpoints do not get used for preemption recovery. eval_steps: List of integers indicating which steps to perform evals. If provided, eval_frequency will be ignored. Performing an eval implies saving a checkpoint that will be used to resume training in the case of preemption. metrics_logger: Used to log all eval metrics during training. See utils.MetricLogger for API definition. init_logger: Used for black box initializers that have learning curves. training_metrics_config: Dict specifying the configuration of the training_metrics_grabber. Set to None to skip logging of advanced training metrics. use_deprecated_checkpointing: Whether to use deprecated checkpointing. Yields: metrics: A dictionary of all eval metrics from the given epoch. """ # NOTE: the initialization RNG should *not* be per-host, as this will create # different sets of weights per host. However, all other RNGs should be # per-host. # TODO(znado,gilmer,gdahl): implement replicating the same initialization # across hosts. rng, init_rng = jax.random.split(rng) rng = jax.random.fold_in(rng, jax.host_id()) rng, data_rng = jax.random.split(rng) # only used if checkpoints_steps is non-empty. checkpoint_dir = os.path.join(train_dir, 'checkpoints') if jax.host_id() == 0: logging.info('Let the training begin!') logging.info('Dataset input shape: %r', hps.input_shape) logging.info('Hyperparameters: %s', hps) if eval_batch_size is None: eval_batch_size = hps.batch_size # Maybe run the initializer. flax_module, batch_stats = initialize(model.flax_module_def, initializer, model.loss_fn, hps.input_shape, hps.output_shape, hps, init_rng, init_logger) if jax.host_id() == 0: utils.log_pytree_shape_and_statistics(flax_module.params) logging.info('train_size: %d,', hps.train_size) lr_fn = schedules.get_schedule_fn(hps.lr_hparams, num_train_steps) optimizer = get_optimizer(hps).create(flax_module) training_metrics_grabber = None if training_metrics_config: training_metrics_grabber = utils.TrainingMetricsGrabber.create( optimizer.target.params, training_metrics_config) (optimizer, batch_stats, training_metrics_grabber, global_step, sum_train_cost, preemption_count, is_restored) = _maybe_restore_latest_checkpoint( unreplicated_optimizer=optimizer, unreplicated_batch_stats=batch_stats, unreplicated_training_metrics_grabber=training_metrics_grabber, train_dir=train_dir, use_deprecated_checkpointing=use_deprecated_checkpointing) if is_restored: preemption_count += 1 # Fold the restored step into the dataset RNG so that we will get a # different shuffle each time we restore, so that we do not repeat a # previous dataset ordering again after restoring. This is not the only # difference in shuffling each pre-emption, because we often times reshuffle # the input files each time in a non-deterministic manner. # # Note that if we are pre-empted more than once per epoch then we will # retrain more on the beginning of the training split, because each time we # restore we refill the shuffle buffer with the first `shuffle_buffer_size` # elements from the training split to continue training. # # Also note that for evaluating on the training split, because we are # reshuffling each time, we will get a new eval_train split each time we are # pre-empted. data_rng = jax.random.fold_in(data_rng, global_step) assert hps.batch_size % (jax.device_count()) == 0 assert eval_batch_size % (jax.device_count()) == 0 dataset = dataset_builder( data_rng, hps.batch_size, eval_batch_size=eval_batch_size, hps=hps, ) # pmap functions for the training loop # in_axes = (optimizer = 0, batch_stats = 0, batch = 0, step = None, # lr = None, rng = None, local_device_index = 0, training_metrics_grabber = 0, # training_metrics_grabber, training_cost ) update_pmapped = jax.pmap(functools.partial( update, training_cost=model.training_cost), axis_name='batch', in_axes=(0, 0, 0, None, None, None, 0, 0)) evaluate_batch_pmapped = jax.pmap(model.evaluate_batch, axis_name='batch') start_time = time.time() start_step = global_step prev_eval_step = start_step def get_step_frequency(cur_step): return float(cur_step - start_step) / (time.time() - start_time) if jax.host_id() == 0: logging.info('Starting training!') # Numpy array of range(0, local_device_count) to send to each device to be # folded into the RNG inside each train step to get a unique per-device RNG. local_device_indices = np.arange(jax.local_device_count()) # Start at the resumed step and continue until we have finished the number of # training steps. If building a dataset iterator using a tf.data.Dataset, in # the case of a batch size that does not evenly divide the training dataset # size, if using `ds.batch(..., drop_remainer=True)` on the training dataset # then the final batch in this iterator will be a partial batch. However, if # `drop_remainer=False`, then this iterator will always return batches of the # same size, and the final batch will have elements from the start of the # (num_epochs + 1)-th epoch. train_iter = itertools.islice(dataset.train_iterator_fn(), global_step, num_train_steps) for batch in train_iter: if global_step in checkpoint_steps and jax.host_id() == 0: save_checkpoint( checkpoint_dir, { 'optimizer': optimizer, 'batch_stats': batch_stats, 'training_metrics_grabber': training_metrics_grabber }, global_step, preemption_count, sum_train_cost, max_to_keep=None, use_deprecated_checkpointing=use_deprecated_checkpointing) batch = data_utils.shard(batch) lr = lr_fn(global_step) optimizer, batch_stats, cost_val, training_metrics_grabber = update_pmapped( optimizer, batch_stats, batch, global_step, lr, rng, local_device_indices, training_metrics_grabber) # Calling float is needed since cost_val is a shape (1,) DeviceArray. sum_train_cost += float(jnp.mean(cost_val)) global_step += 1 # TODO(gdahl, gilmer): consider moving this test up. # NB: Since this test is after we increment global_step, having 0 in # eval_steps does nothing. if should_eval(global_step, eval_frequency, eval_steps): batch_stats = _maybe_sync_batchnorm_stats(batch_stats) report, eval_time = eval_metrics(optimizer.target, batch_stats, dataset, eval_num_batches, eval_train_num_batches, evaluate_batch_pmapped) mean_train_cost = sum_train_cost / max( 1, global_step - prev_eval_step) report.update(learning_rate=float(lr), global_step=global_step, epoch=global_step * hps.batch_size // hps.train_size, steps_per_sec=get_step_frequency(global_step), eval_time=eval_time, preemption_count=preemption_count, train_cost=mean_train_cost) yield report if jax.host_id() == 0: _log_epoch_report(report, metrics_logger) _maybe_log_training_metrics(training_metrics_grabber, metrics_logger) save_checkpoint( train_dir, { 'optimizer': optimizer, 'batch_stats': batch_stats, 'training_metrics_grabber': training_metrics_grabber }, global_step, preemption_count, sum_train_cost, use_deprecated_checkpointing=use_deprecated_checkpointing) sum_train_cost = 0.0 prev_eval_step = global_step # Always log and checkpoint on host 0 at the end of training. # If we moved where in the loop body evals happen then we would not need this # test. if prev_eval_step != num_train_steps: batch_stats = _maybe_sync_batchnorm_stats(batch_stats) report, eval_time = eval_metrics(optimizer.target, batch_stats, dataset, eval_num_batches, eval_train_num_batches, evaluate_batch_pmapped) lr = lr_fn(global_step) # Correct the average for the final partial epoch. mean_train_cost = sum_train_cost / max(1, global_step - prev_eval_step) report.update(learning_rate=float(lr), global_step=global_step, epoch=global_step * hps.batch_size // hps.train_size, steps_per_sec=get_step_frequency(global_step), eval_time=eval_time, preemption_count=preemption_count, train_cost=mean_train_cost) yield report if jax.host_id() == 0: _log_epoch_report(report, metrics_logger) _maybe_log_training_metrics(training_metrics_grabber, metrics_logger) save_checkpoint( train_dir, { 'optimizer': optimizer, 'batch_stats': batch_stats, 'training_metrics_grabber': training_metrics_grabber }, global_step, preemption_count, sum_train_cost, use_deprecated_checkpointing=use_deprecated_checkpointing) # To make sure the last checkpoint was correctly saved. checkpoint.wait_for_checkpoint_save()