def test_get_strategy_scope(self): ds = distribute_utils.get_distribution_strategy('one_device', num_gpus=0) with distribute_utils.get_strategy_scope(ds): self.assertIs(tf.distribute.get_strategy(), ds) with distribute_utils.get_strategy_scope(None): self.assertIsNot(tf.distribute.get_strategy(), ds)
def train(self, train_data, validation_data, epochs=10, batch_size=4): if len(train_data) < batch_size: raise ValueError( 'The size of the train_data (%d) couldn\'t be smaller ' 'than batch_size (%d). To solve this problem, set ' 'the batch_size smaller or increase the size of the ' 'train_data.' % (len(train_data), batch_size)) with distribute_utils.get_strategy_scope(self.model_spec.strategy): train_input_fn, train_steps = self._get_input_fn_and_steps( train_data, batch_size, is_training=True) validation_input_fn, validation_steps = self._get_input_fn_and_steps( validation_data, batch_size, is_training=False) # TODO(wangtz): This should probably be fused into _get_input_fn_and_steps train_ds = _get_dataset(train_input_fn) validation_ds = _get_dataset(validation_input_fn) self.model = self.model_spec.create_model() return self.model_spec.run_classifier( self.model, epochs, train_ds, train_steps, validation_ds, validation_steps, callbacks=self._keras_callbacks(self.model_spec.model_dir))
def run_classifier(self, train_input_fn, validation_input_fn, epochs, steps_per_epoch, validation_steps, num_classes): """Creates classifier and runs the classifier training.""" warmup_steps = int(epochs * steps_per_epoch * 0.1) initial_lr = self.learning_rate with distribute_utils.get_strategy_scope(self.strategy): training_dataset = train_input_fn() evaluation_dataset = None if validation_input_fn is not None: evaluation_dataset = validation_input_fn() optimizer = optimization.create_optimizer(initial_lr, steps_per_epoch * epochs, warmup_steps) bert_model = self.create_model(num_classes, optimizer) summary_dir = os.path.join(self.model_dir, 'summaries') summary_callback = tf.keras.callbacks.TensorBoard(summary_dir) checkpoint_path = os.path.join(self.model_dir, 'checkpoint') checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( checkpoint_path, save_weights_only=True) bert_model.fit(x=training_dataset, validation_data=evaluation_dataset, steps_per_epoch=steps_per_epoch, epochs=epochs, validation_steps=validation_steps, callbacks=[summary_callback, checkpoint_callback]) return bert_model
def train(self, train_input_fn, epochs, steps_per_epoch, **kwargs): """Run bert QA training.""" warmup_steps = int(epochs * steps_per_epoch * 0.1) def _loss_fn(positions, logits): """Get losss function for QA model.""" loss = tf.keras.losses.sparse_categorical_crossentropy( positions, logits, from_logits=True) return tf.reduce_mean(loss) with distribute_utils.get_strategy_scope(self.strategy): training_dataset = train_input_fn() bert_model = self.create_model() optimizer = optimization.create_optimizer(self.learning_rate, steps_per_epoch * epochs, warmup_steps) bert_model.compile( optimizer=optimizer, loss=_loss_fn, loss_weights=[0.5, 0.5]) if not bert_model.trainable_variables: tf.compat.v1.logging.warning( 'Trainable variables in the model are empty.') return bert_model bert_model.fit( x=training_dataset, steps_per_epoch=steps_per_epoch, epochs=epochs, **kwargs) return bert_model
def run_classifier(self, train_input_fn, validation_input_fn, epochs, steps_per_epoch, validation_steps, num_classes, **kwargs): """Creates classifier and runs the classifier training.""" warmup_steps = int(epochs * steps_per_epoch * 0.1) initial_lr = self.learning_rate with distribute_utils.get_strategy_scope(self.strategy): training_dataset = train_input_fn() evaluation_dataset = None if validation_input_fn is not None: evaluation_dataset = validation_input_fn() optimizer = optimization.create_optimizer(initial_lr, steps_per_epoch * epochs, warmup_steps) bert_model = self.create_model(num_classes, optimizer) bert_model.fit( x=training_dataset, validation_data=evaluation_dataset, steps_per_epoch=steps_per_epoch, epochs=epochs, validation_steps=validation_steps, **kwargs) return bert_model
def train(self, train_ds, epochs, steps_per_epoch, **kwargs): """Run bert QA training. Args: train_ds: tf.data.Dataset, training data to be fed in tf.keras.Model.fit(). epochs: Integer, training epochs. steps_per_epoch: Integer or None. Total number of steps (batches of samples) before declaring one epoch finished and starting the next epoch. If `steps_per_epoch` is None, the epoch will run until the input dataset is exhausted. **kwargs: Other parameters used in the tf.keras.Model.fit(). Returns: tf.keras.Model, the keras model that's already trained. """ if steps_per_epoch is None: logging.info( 'steps_per_epoch is None, use %d as the estimated steps_per_epoch', model_util.ESTIMITED_STEPS_PER_EPOCH) steps_per_epoch = model_util.ESTIMITED_STEPS_PER_EPOCH total_steps = steps_per_epoch * epochs warmup_steps = int(total_steps * 0.1) def _loss_fn(positions, logits): """Get losss function for QA model.""" loss = tf.keras.losses.sparse_categorical_crossentropy( positions, logits, from_logits=True) return tf.reduce_mean(loss) with distribute_utils.get_strategy_scope(self.strategy): bert_model = self.create_model() optimizer = optimization.create_optimizer(self.learning_rate, total_steps, warmup_steps) bert_model.compile(optimizer=optimizer, loss=_loss_fn, loss_weights=[0.5, 0.5]) if not bert_model.trainable_variables: tf.compat.v1.logging.warning( 'Trainable variables in the model are empty.') return bert_model bert_model.fit(x=train_ds, epochs=epochs, **kwargs) return bert_model
def eval(self): """Evaluates the model.""" distribution_strategy = self.distribution_strategy if self.use_tpu else None # We only want to create the model under DS scope for TPU case. # When 'distribution_strategy' is None, a no-op DummyContextManager will # be used. with distribute_utils.get_strategy_scope(distribution_strategy): if not self.predict_model: self.predict_model = transformer.create_model(self.params, False) self._load_weights_if_possible( self.predict_model, tf.train.latest_checkpoint(self.flags_obj.model_dir)) self.predict_model.summary() return evaluate_and_log_bleu( self.predict_model, self.params, self.flags_obj.bleu_source, self.flags_obj.bleu_ref, self.flags_obj.vocab_file, distribution_strategy)
def run_classifier(self, train_ds, validation_ds, epochs, steps_per_epoch, num_classes, **kwargs): """Creates classifier and runs the classifier training. Args: train_ds: tf.data.Dataset, training data to be fed in tf.keras.Model.fit(). validation_ds: tf.data.Dataset, validation data to be fed in tf.keras.Model.fit(). epochs: Integer, training epochs. steps_per_epoch: Integer or None. Total number of steps (batches of samples) before declaring one epoch finished and starting the next epoch. If `steps_per_epoch` is None, the epoch will run until the input dataset is exhausted. num_classes: Interger, number of classes. **kwargs: Other parameters used in the tf.keras.Model.fit(). Returns: tf.keras.Model, the keras model that's already trained. """ if steps_per_epoch is None: logging.info( 'steps_per_epoch is None, use %d as the estimated steps_per_epoch', model_util.ESTIMITED_STEPS_PER_EPOCH) steps_per_epoch = model_util.ESTIMITED_STEPS_PER_EPOCH total_steps = steps_per_epoch * epochs warmup_steps = int(total_steps * 0.1) initial_lr = self.learning_rate with distribute_utils.get_strategy_scope(self.strategy): optimizer = optimization.create_optimizer(initial_lr, total_steps, warmup_steps) bert_model = self.create_model(num_classes, optimizer) for i in range(epochs): bert_model.fit( x=train_ds, initial_epoch=i, epochs=i + 1, validation_data=validation_ds, **kwargs) return bert_model
def train(self, train_data, validation_data, epochs=10, batch_size=4): if len(train_data) < batch_size: raise ValueError('The size of the train_data (%d) couldn\'t be smaller ' 'than batch_size (%d). To solve this problem, set ' 'the batch_size smaller or increase the size of the ' 'train_data.' % (len(train_data), batch_size)) with distribute_utils.get_strategy_scope(self.model_spec.strategy): train_ds = train_data.gen_dataset(batch_size, is_training=True) train_steps = len(train_data) // batch_size validation_ds = validation_data.gen_dataset( batch_size, is_training=False) if validation_data else None self.model = self.model_spec.create_model() return self.model_spec.run_classifier( self.model, epochs, train_ds, train_steps, validation_ds, callbacks=self._keras_callbacks(self.model_spec.model_dir))
def train(self, train_input_fn, epochs, steps_per_epoch): """Run bert QA training.""" warmup_steps = int(epochs * steps_per_epoch * 0.1) def _loss_fn(positions, logits): """Get losss function for QA model.""" loss = tf.keras.losses.sparse_categorical_crossentropy( positions, logits, from_logits=True) return tf.reduce_mean(loss) with distribute_utils.get_strategy_scope(self.strategy): training_dataset = train_input_fn() bert_model = self.create_model() optimizer = optimization.create_optimizer(self.learning_rate, steps_per_epoch * epochs, warmup_steps) bert_model.compile(optimizer=optimizer, loss=_loss_fn, loss_weights=[0.5, 0.5]) summary_dir = os.path.join(self.model_dir, 'summaries') summary_callback = tf.keras.callbacks.TensorBoard(summary_dir) checkpoint_path = os.path.join(self.model_dir, 'checkpoint') checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( checkpoint_path, save_weights_only=True) if not bert_model.trainable_variables: tf.compat.v1.logging.warning( 'Trainable variables in the model are empty.') return bert_model bert_model.fit(x=training_dataset, steps_per_epoch=steps_per_epoch, epochs=epochs, callbacks=[summary_callback, checkpoint_callback]) return bert_model
def train(self): """Trains the model.""" params = self.params flags_obj = self.flags_obj # Sets config options. keras_utils.set_session_config(enable_xla=flags_obj.enable_xla) _ensure_dir(flags_obj.model_dir) with distribute_utils.get_strategy_scope(self.distribution_strategy): model = transformer.create_model(params, is_train=True) opt = self._create_optimizer() current_step = 0 checkpoint = tf.train.Checkpoint(model=model, optimizer=opt) latest_checkpoint = tf.train.latest_checkpoint(flags_obj.model_dir) if latest_checkpoint: checkpoint.restore(latest_checkpoint) logging.info("Loaded checkpoint %s", latest_checkpoint) current_step = opt.iterations.numpy() if params["use_ctl"]: train_loss_metric = tf.keras.metrics.Mean( "training_loss", dtype=tf.float32) if params["enable_tensorboard"]: summary_writer = tf.summary.create_file_writer( os.path.join(flags_obj.model_dir, "summary")) else: summary_writer = tf.summary.create_noop_writer() train_metrics = [train_loss_metric] if params["enable_metrics_in_training"]: train_metrics = train_metrics + model.metrics else: model.compile(opt) model.summary() if self.use_tpu: # Different from experimental_distribute_dataset, # distribute_datasets_from_function requires # per-replica/local batch size. params["batch_size"] /= self.distribution_strategy.num_replicas_in_sync train_ds = ( self.distribution_strategy.distribute_datasets_from_function( lambda ctx: data_pipeline.train_input_fn(params, ctx))) else: train_ds = data_pipeline.train_input_fn(params) map_data_fn = data_pipeline.map_data_for_transformer_fn train_ds = train_ds.map( map_data_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) if params["use_ctl"]: train_ds_iterator = iter(train_ds) callbacks = self._create_callbacks(flags_obj.model_dir, params) # Only TimeHistory callback is supported for CTL if params["use_ctl"]: callbacks = [cb for cb in callbacks if isinstance(cb, keras_utils.TimeHistory)] @tf.function def train_steps(iterator, steps): """Training steps function for TPU runs. Args: iterator: The input iterator of the training dataset. steps: An integer, the number of training steps. Returns: A float, the loss value. """ def _step_fn(inputs): """Per-replica step function.""" inputs, targets = inputs with tf.GradientTape() as tape: logits = model([inputs, targets], training=True) loss = metrics.transformer_loss(logits, targets, params["label_smoothing"], params["vocab_size"]) # Scales the loss, which results in using the average loss across all # of the replicas for backprop. scaled_loss = loss / self.distribution_strategy.num_replicas_in_sync # De-dupes variables due to keras tracking issues. tvars = list({id(v): v for v in model.trainable_variables}.values()) grads = tape.gradient(scaled_loss, tvars) opt.apply_gradients(zip(grads, tvars)) # For reporting, the metric takes the mean of losses. train_loss_metric.update_state(loss) for _ in tf.range(steps): train_loss_metric.reset_states() self.distribution_strategy.run( _step_fn, args=(next(iterator),)) cased_score, uncased_score = None, None cased_score_history, uncased_score_history = [], [] while current_step < flags_obj.train_steps: remaining_steps = flags_obj.train_steps - current_step train_steps_per_eval = ( remaining_steps if remaining_steps < flags_obj.steps_between_evals else flags_obj.steps_between_evals) current_iteration = current_step // flags_obj.steps_between_evals logging.info( "Start train iteration at global step:{}".format(current_step)) history = None if params["use_ctl"]: if not self.use_tpu: raise NotImplementedError( "Custom training loop on GPUs is not implemented.") # Runs training steps. with summary_writer.as_default(): for cb in callbacks: cb.on_epoch_begin(current_iteration) cb.on_batch_begin(0) train_steps( train_ds_iterator, tf.convert_to_tensor(train_steps_per_eval, dtype=tf.int32)) current_step += train_steps_per_eval train_loss = train_loss_metric.result().numpy().astype(float) logging.info("Train Step: %d/%d / loss = %s", current_step, flags_obj.train_steps, train_loss) for cb in callbacks: cb.on_batch_end(train_steps_per_eval - 1) cb.on_epoch_end(current_iteration) if params["enable_tensorboard"]: for metric_obj in train_metrics: tf.summary.scalar(metric_obj.name, metric_obj.result(), current_step) summary_writer.flush() for cb in callbacks: cb.on_train_end() if flags_obj.enable_checkpointing: # avoid check-pointing when running for benchmarking. checkpoint_name = checkpoint.save( os.path.join(flags_obj.model_dir, "ctl_step_{}.ckpt".format(current_step))) logging.info("Saved checkpoint to %s", checkpoint_name) else: if self.use_tpu: raise NotImplementedError( "Keras model.fit on TPUs is not implemented.") history = model.fit( train_ds, initial_epoch=current_iteration, epochs=current_iteration + 1, steps_per_epoch=train_steps_per_eval, callbacks=callbacks, # If TimeHistory is enabled, progress bar would be messy. Increase # the verbose level to get rid of it. verbose=(2 if flags_obj.enable_time_history else 1)) current_step += train_steps_per_eval logging.info("Train history: {}".format(history.history)) logging.info("End train iteration at global step:{}".format(current_step)) if (flags_obj.bleu_source and flags_obj.bleu_ref): uncased_score, cased_score = self.eval() cased_score_history.append([current_iteration + 1, cased_score]) uncased_score_history.append([current_iteration + 1, uncased_score]) stats = ({ "loss": train_loss } if history is None else {}) misc.update_stats(history, stats, callbacks) if uncased_score and cased_score: stats["bleu_uncased"] = uncased_score stats["bleu_cased"] = cased_score stats["bleu_uncased_history"] = uncased_score_history stats["bleu_cased_history"] = cased_score_history return stats
def run_ncf(_): """Run NCF training and eval with Keras.""" keras_utils.set_session_config(enable_xla=FLAGS.enable_xla) if FLAGS.seed is not None: print("Setting tf seed") tf.random.set_seed(FLAGS.seed) model_helpers.apply_clean(FLAGS) if FLAGS.dtype == "fp16" and FLAGS.fp16_implementation == "keras": tf.keras.mixed_precision.set_global_policy("mixed_float16") strategy = distribute_utils.get_distribution_strategy( distribution_strategy=FLAGS.distribution_strategy, num_gpus=FLAGS.num_gpus, tpu_address=FLAGS.tpu) params = ncf_common.parse_flags(FLAGS) params["distribute_strategy"] = strategy params["use_tpu"] = (FLAGS.distribution_strategy == "tpu") if params["use_tpu"] and not params["keras_use_ctl"]: logging.error( "Custom training loop must be used when using TPUStrategy.") return batch_size = params["batch_size"] time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps) callbacks = [time_callback] producer, input_meta_data = None, None generate_input_online = params["train_dataset_path"] is None if generate_input_online: # Start data producing thread. num_users, num_items, _, _, producer = ncf_common.get_inputs(params) producer.start() per_epoch_callback = IncrementEpochCallback(producer) callbacks.append(per_epoch_callback) else: assert params["eval_dataset_path"] and params["input_meta_data_path"] with tf.io.gfile.GFile(params["input_meta_data_path"], "rb") as reader: input_meta_data = json.loads(reader.read().decode("utf-8")) num_users = input_meta_data["num_users"] num_items = input_meta_data["num_items"] params["num_users"], params["num_items"] = num_users, num_items if FLAGS.early_stopping: early_stopping_callback = CustomEarlyStopping( "val_HR_METRIC", desired_value=FLAGS.hr_threshold) callbacks.append(early_stopping_callback) (train_input_dataset, eval_input_dataset, num_train_steps, num_eval_steps) = \ (ncf_input_pipeline.create_ncf_input_data( params, producer, input_meta_data, strategy)) steps_per_epoch = None if generate_input_online else num_train_steps with distribute_utils.get_strategy_scope(strategy): keras_model = _get_keras_model(params) optimizer = tf.keras.optimizers.Adam( learning_rate=params["learning_rate"], beta_1=params["beta1"], beta_2=params["beta2"], epsilon=params["epsilon"]) if FLAGS.fp16_implementation == "graph_rewrite": optimizer = \ tf.compat.v1.train.experimental.enable_mixed_precision_graph_rewrite( optimizer, loss_scale=flags_core.get_loss_scale(FLAGS, default_for_fp16="dynamic")) elif FLAGS.dtype == "fp16": loss_scale = flags_core.get_loss_scale(FLAGS, default_for_fp16="dynamic") # Note Model.compile automatically wraps the optimizer with a # LossScaleOptimizer using dynamic loss scaling. We explicitly wrap it # here for the case where a custom training loop or fixed loss scale is # used. if loss_scale == "dynamic": optimizer = tf.keras.mixed_precision.LossScaleOptimizer( optimizer) else: optimizer = tf.keras.mixed_precision.LossScaleOptimizer( optimizer, dynamic=False, initial_scale=loss_scale) if params["keras_use_ctl"]: train_loss, eval_results = run_ncf_custom_training( params, strategy, keras_model, optimizer, callbacks, train_input_dataset, eval_input_dataset, num_train_steps, num_eval_steps, generate_input_online=generate_input_online) else: keras_model.compile(optimizer=optimizer, run_eagerly=FLAGS.run_eagerly) if not FLAGS.ml_perf: # Create Tensorboard summary and checkpoint callbacks. summary_dir = os.path.join(FLAGS.model_dir, "summaries") summary_callback = tf.keras.callbacks.TensorBoard(summary_dir) checkpoint_path = os.path.join(FLAGS.model_dir, "checkpoint") checkpoint_callback = tf.keras.callbacks.ModelCheckpoint( checkpoint_path, save_weights_only=True) callbacks += [summary_callback, checkpoint_callback] history = keras_model.fit(train_input_dataset, epochs=FLAGS.train_epochs, steps_per_epoch=steps_per_epoch, callbacks=callbacks, validation_data=eval_input_dataset, validation_steps=num_eval_steps, verbose=2) logging.info("Training done. Start evaluating") eval_loss_and_metrics = keras_model.evaluate(eval_input_dataset, steps=num_eval_steps, verbose=2) logging.info("Keras evaluation is done.") # Keras evaluate() API returns scalar loss and metric values from # evaluation as a list. Here, the returned list would contain # [evaluation loss, hr sum, hr count]. eval_hit_rate = eval_loss_and_metrics[1] / eval_loss_and_metrics[2] # Format evaluation result into [eval loss, eval hit accuracy]. eval_results = [eval_loss_and_metrics[0], eval_hit_rate] if history and history.history: train_history = history.history train_loss = train_history["loss"][-1] stats = build_stats(train_loss, eval_results, time_callback) return stats
def train_and_eval( params: base_configs.ExperimentConfig, strategy_override: tf.distribute.Strategy) -> Mapping[str, Any]: """Runs the train and eval path using compile/fit.""" logging.info('Running train and eval.') distribute_utils.configure_cluster(params.runtime.worker_hosts, params.runtime.task_index) # Note: for TPUs, strategy and scope should be created before the dataset strategy = strategy_override or distribute_utils.get_distribution_strategy( distribution_strategy=params.runtime.distribution_strategy, all_reduce_alg=params.runtime.all_reduce_alg, num_gpus=params.runtime.num_gpus, tpu_address=params.runtime.tpu) strategy_scope = distribute_utils.get_strategy_scope(strategy) logging.info('Detected %d devices.', strategy.num_replicas_in_sync if strategy else 1) label_smoothing = params.model.loss.label_smoothing one_hot = label_smoothing and label_smoothing > 0 builders = _get_dataset_builders(params, strategy, one_hot) datasets = [ builder.build(strategy) if builder else None for builder in builders ] # Unpack datasets and builders based on train/val/test splits train_builder, validation_builder = builders # pylint: disable=unbalanced-tuple-unpacking train_dataset, validation_dataset = datasets train_epochs = params.train.epochs train_steps = params.train.steps or train_builder.num_steps validation_steps = params.evaluation.steps or validation_builder.num_steps initialize(params, train_builder) logging.info('Global batch size: %d', train_builder.global_batch_size) with strategy_scope: model_params = params.model.model_params.as_dict() model = get_models()[params.model.name](**model_params) learning_rate = optimizer_factory.build_learning_rate( params=params.model.learning_rate, batch_size=train_builder.global_batch_size, train_epochs=train_epochs, train_steps=train_steps) optimizer = optimizer_factory.build_optimizer( optimizer_name=params.model.optimizer.name, base_learning_rate=learning_rate, params=params.model.optimizer.as_dict(), model=model) optimizer = performance.configure_optimizer( optimizer, use_float16=train_builder.dtype == 'float16', loss_scale=get_loss_scale(params)) metrics_map = _get_metrics(one_hot) metrics = [metrics_map[metric] for metric in params.train.metrics] steps_per_loop = train_steps if params.train.set_epoch_loop else 1 if one_hot: loss_obj = tf.keras.losses.CategoricalCrossentropy( label_smoothing=params.model.loss.label_smoothing) else: loss_obj = tf.keras.losses.SparseCategoricalCrossentropy() model.compile( optimizer=optimizer, loss=loss_obj, metrics=metrics, steps_per_execution=steps_per_loop) initial_epoch = 0 if params.train.resume_checkpoint: initial_epoch = resume_from_checkpoint( model=model, model_dir=params.model_dir, train_steps=train_steps) callbacks = custom_callbacks.get_callbacks( model_checkpoint=params.train.callbacks.enable_checkpoint_and_export, include_tensorboard=params.train.callbacks.enable_tensorboard, time_history=params.train.callbacks.enable_time_history, track_lr=params.train.tensorboard.track_lr, write_model_weights=params.train.tensorboard.write_model_weights, initial_step=initial_epoch * train_steps, batch_size=train_builder.global_batch_size, log_steps=params.train.time_history.log_steps, model_dir=params.model_dir, backup_and_restore=params.train.callbacks.enable_backup_and_restore) serialize_config(params=params, model_dir=params.model_dir) if params.evaluation.skip_eval: validation_kwargs = {} else: validation_kwargs = { 'validation_data': validation_dataset, 'validation_steps': validation_steps, 'validation_freq': params.evaluation.epochs_between_evals, } history = model.fit( train_dataset, epochs=train_epochs, steps_per_epoch=train_steps, initial_epoch=initial_epoch, callbacks=callbacks, verbose=2, **validation_kwargs) validation_output = None if not params.evaluation.skip_eval: validation_output = model.evaluate( validation_dataset, steps=validation_steps, verbose=2) # TODO(dankondratyuk): eval and save final test accuracy stats = common.build_stats(history, validation_output, callbacks) return stats
def run_customized_training_loop( # pylint: disable=invalid-name _sentinel=None, # pylint: enable=invalid-name strategy=None, model_fn=None, loss_fn=None, scale_loss=True, model_dir=None, train_input_fn=None, steps_per_epoch=None, num_eval_per_epoch=1, steps_per_loop=None, epochs=1, eval_input_fn=None, eval_steps=None, metric_fn=None, init_checkpoint=None, custom_callbacks=None, run_eagerly=False, sub_model_export_name=None, explicit_allreduce=False, pre_allreduce_callbacks=None, post_allreduce_callbacks=None, train_summary_interval=0, allreduce_bytes_per_pack=0): """Run BERT pretrain model training using low-level API. Arguments: _sentinel: Used to prevent positional parameters. Internal, do not use. strategy: Distribution strategy on which to run low level training loop. model_fn: Function that returns a tuple (model, sub_model). Caller of this function should add optimizer to the `model` via calling `model.compile()` API or manually setting `model.optimizer` attribute. Second element of the returned tuple(sub_model) is an optional sub model to be used for initial checkpoint -- if provided. loss_fn: Function with signature func(labels, logits) and returns a loss tensor. scale_loss: Whether to divide the raw loss by number of replicas before gradients calculation. model_dir: Model directory used during training for restoring/saving model weights. train_input_fn: Function that returns a tf.data.Dataset used for training. steps_per_epoch: Number of steps to run per epoch. At the end of each epoch, model checkpoint will be saved and evaluation will be conducted if evaluation dataset is provided. num_eval_per_epoch: Number of evaluations per epoch. steps_per_loop: Number of steps per graph-mode loop. In order to reduce communication in eager context, training logs are printed every steps_per_loop. epochs: Number of epochs to train. eval_input_fn: Function that returns evaluation dataset. If none, evaluation is skipped. eval_steps: Number of steps to run evaluation. Required if `eval_input_fn` is not none. metric_fn: A metrics function that returns either a Keras Metric object or a list of Keras Metric objects to record evaluation result using evaluation dataset or with training dataset after every epoch. init_checkpoint: Optional checkpoint to load to `sub_model` returned by `model_fn`. custom_callbacks: A list of Keras Callbacks objects to run during training. More specifically, `on_train_begin(), on_train_end(), on_batch_begin()`, `on_batch_end()`, `on_epoch_begin()`, `on_epoch_end()` methods are invoked during training. Note that some metrics may be missing from `logs`. run_eagerly: Whether to run model training in pure eager execution. This should be disable for TPUStrategy. sub_model_export_name: If not None, will export `sub_model` returned by `model_fn` into checkpoint files. The name of intermediate checkpoint file is {sub_model_export_name}_step_{step}.ckpt and the last checkpint's name is {sub_model_export_name}.ckpt; if None, `sub_model` will not be exported as checkpoint. explicit_allreduce: Whether to explicitly perform gradient allreduce, instead of relying on implicit allreduce in optimizer.apply_gradients(). default is False. For now, if training using FP16 mixed precision, explicit allreduce will aggregate gradients in FP16 format. For TPU and GPU training using FP32, explicit allreduce will aggregate gradients in FP32 format. pre_allreduce_callbacks: A list of callback functions that takes gradients and model variables pairs as input, manipulate them, and returns a new gradients and model variables paris. The callback functions will be invoked in the list order and before gradients are allreduced. With mixed precision training, the pre_allreduce_allbacks will be applied on scaled_gradients. Default is no callbacks. Only used when explicit_allreduce=True. post_allreduce_callbacks: A list of callback functions that takes gradients and model variables pairs as input, manipulate them, and returns a new gradients and model variables paris. The callback functions will be invoked in the list order and right before gradients are applied to variables for updates. Default is no callbacks. Only used when explicit_allreduce=True. train_summary_interval: Step interval for training summaries. If the value is a negative number, then training summaries are not enabled. allreduce_bytes_per_pack: A non-negative integer. Breaks collective operations into packs of certain size. If it's zero, all gradients are in one pack. Breaking gradient into packs could enable overlap between allreduce and backprop computation. This flag only takes effect when explicit_allreduce is set to True.' Returns: Trained model. Raises: ValueError: (1) When model returned by `model_fn` does not have optimizer attribute or when required parameters are set to none. (2) eval args are not specified correctly. (3) metric_fn must be a callable if specified. (4) sub_model_checkpoint_name is specified, but `sub_model` returned by `model_fn` is None. """ if _sentinel is not None: raise ValueError('only call `run_customized_training_loop()` ' 'with named arguments.') required_arguments = [ strategy, model_fn, loss_fn, model_dir, steps_per_epoch, train_input_fn ] steps_between_evals = int(steps_per_epoch / num_eval_per_epoch) if [arg for arg in required_arguments if arg is None]: raise ValueError('`strategy`, `model_fn`, `loss_fn`, `model_dir`, ' '`steps_per_epoch` and `train_input_fn` are required ' 'parameters.') if not steps_per_loop: if tf.config.list_logical_devices('TPU'): # One can't fully utilize a TPU with steps_per_loop=1, so in this case # default users to a more useful value. steps_per_loop = min(1000, steps_between_evals) else: steps_per_loop = 1 logging.info('steps_per_loop not specified. Using steps_per_loop=%d', steps_per_loop) if steps_per_loop > steps_between_evals: logging.warning( 'steps_per_loop: %d is specified to be greater than ' ' steps_between_evals: %d, we will use steps_between_evals as' ' steps_per_loop.', steps_per_loop, steps_between_evals) steps_per_loop = steps_between_evals assert tf.executing_eagerly() if run_eagerly: if isinstance(strategy, tf.distribute.experimental.TPUStrategy): raise ValueError( 'TPUStrategy should not run eagerly as it heavily relies on graph' ' optimization for the distributed system.') if eval_input_fn and eval_steps is None: raise ValueError( '`eval_step` is required when `eval_input_fn ` is not none.') if metric_fn and not callable(metric_fn): raise ValueError( 'if `metric_fn` is specified, metric_fn must be a callable.') total_training_steps = steps_per_epoch * epochs train_iterator = _get_input_iterator(train_input_fn, strategy) eval_loss_metric = tf.keras.metrics.Mean('training_loss', dtype=tf.float32) with distribute_utils.get_strategy_scope(strategy): # To correctly place the model weights on accelerators, # model and optimizer should be created in scope. model, sub_model = model_fn() if not hasattr(model, 'optimizer'): raise ValueError('User should set optimizer attribute to model ' 'inside `model_fn`.') if sub_model_export_name and sub_model is None: raise ValueError('sub_model_export_name is specified as %s, but ' 'sub_model is None.' % sub_model_export_name) callback_list = tf.keras.callbacks.CallbackList( callbacks=custom_callbacks, model=model) optimizer = model.optimizer if init_checkpoint: logging.info( 'Checkpoint file %s found and restoring from ' 'initial checkpoint for core model.', init_checkpoint) checkpoint = tf.train.Checkpoint(model=sub_model) checkpoint.restore(init_checkpoint).assert_existing_objects_matched() logging.info('Loading from checkpoint file completed') train_loss_metric = tf.keras.metrics.Mean('training_loss', dtype=tf.float32) eval_metrics = metric_fn() if metric_fn else [] if not isinstance(eval_metrics, list): eval_metrics = [eval_metrics] # If evaluation is required, make a copy of metric as it will be used by # both train and evaluation. train_metrics = [ metric.__class__.from_config(metric.get_config()) for metric in eval_metrics ] # Create summary writers if _should_export_summary(strategy): summary_dir = os.path.join(model_dir, 'summaries') else: # In multi worker training we need every worker to write summary, because # variables can trigger synchronization on read and synchronization needs # all workers to participate. summary_dir = tempfile.mkdtemp() eval_summary_writer = tf.summary.create_file_writer( os.path.join(summary_dir, 'eval')) last_summary_step = 0 if steps_per_loop >= _MIN_SUMMARY_STEPS and train_summary_interval >= 0: # Only writes summary when the stats are collected sufficiently over # enough steps. train_summary_writer = tf.summary.create_file_writer( os.path.join(summary_dir, 'train')) else: train_summary_writer = tf.summary.create_noop_writer() # Collects training variables. training_vars = model.trainable_variables def _replicated_step(inputs): """Replicated training step.""" inputs, labels = inputs with tf.GradientTape() as tape: model_outputs = model(inputs, training=True) loss = loss_fn(labels, model_outputs) # Raw loss is used for reporting in metrics/logs. raw_loss = loss if scale_loss: # Scales down the loss for gradients to be invariant from replicas. loss = loss / strategy.num_replicas_in_sync if explicit_allreduce: grad_utils.minimize_using_explicit_allreduce(tape, optimizer, loss, training_vars, pre_allreduce_callbacks, post_allreduce_callbacks, allreduce_bytes_per_pack) else: if isinstance(optimizer, tf.keras.mixed_precision.LossScaleOptimizer): with tape: scaled_loss = optimizer.get_scaled_loss(loss) scaled_grads = tape.gradient(scaled_loss, training_vars) grads = optimizer.get_unscaled_gradients(scaled_grads) else: grads = tape.gradient(loss, training_vars) optimizer.apply_gradients(zip(grads, training_vars)) # For reporting, the metric takes the mean of losses. train_loss_metric.update_state(raw_loss) for metric in train_metrics: metric.update_state(labels, model_outputs) @tf.function def train_steps(iterator, steps): """Performs distributed training steps in a loop. Args: iterator: the distributed iterator of training datasets. steps: an tf.int32 integer tensor to specify number of steps to run inside host training loop. Raises: ValueError: Any of the arguments or tensor shapes are invalid. """ if not isinstance(steps, tf.Tensor): raise ValueError('steps should be an Tensor. Python object may cause ' 'retracing.') for _ in tf.range(steps): strategy.run(_replicated_step, args=(next(iterator),)) def train_single_step(iterator): """Performs a distributed training step. Args: iterator: the distributed iterator of training datasets. Raises: ValueError: Any of the arguments or tensor shapes are invalid. """ strategy.run(_replicated_step, args=(next(iterator),)) def test_step(iterator): """Calculates evaluation metrics on distributed devices.""" def _test_step_fn(inputs): """Replicated accuracy calculation.""" inputs, labels = inputs model_outputs = model(inputs, training=False) for metric in eval_metrics: metric.update_state(labels, model_outputs) return model_outputs, labels outputs, labels = strategy.run(_test_step_fn, args=(next(iterator),)) outputs = tf.nest.map_structure(strategy.experimental_local_results, outputs) labels = tf.nest.map_structure(strategy.experimental_local_results, labels) return outputs, labels if not run_eagerly: train_single_step = tf.function(train_single_step) test_step = tf.function(test_step) def _run_evaluation(current_training_step, test_iterator): """Runs validation steps and aggregate metrics. Args: current_training_step: tf.int32 tensor containing the current step. test_iterator: distributed iterator of test datasets. Returns: A dict of metic names and values. """ # The last batch of the evaluation is often smaller than previous ones. # Moreover, in some distributed pieces it might even be empty. Therefore, # different from the way training_loss is calculated, it is needed to # gather all the logits and labels here to calculate the evaluation loss # outside. loss_list, loss_weights = list(), list() for _ in range(eval_steps): outputs, labels = test_step(test_iterator) for cur_logits, cur_labels in zip(outputs, labels): # This is to handle cases when cur_labels is not a single tensor, # but a dict of tensors. cur_weight = tf.shape(tf.nest.flatten(cur_labels)[0])[0] if cur_weight != 0: loss_list.append(loss_fn(cur_labels, cur_logits).numpy()) loss_weights.append(cur_weight) # The sample_weights are the actual number of examples in each batch, # a summation of numbers of examples in each replica if using # distributed training. eval_loss_metric.update_state(loss_list, sample_weight=loss_weights) logs = {} with eval_summary_writer.as_default(): for metric in [eval_loss_metric] + eval_metrics + model.metrics: metric_value = _float_metric_value(metric) logs[metric.name] = metric_value logging.info('Step: [%d] Validation %s = %f', current_training_step, metric.name, metric_value) tf.summary.scalar( metric.name, metric_value, step=current_training_step) eval_summary_writer.flush() return logs # Training loop starts here. checkpoint = tf.train.Checkpoint( model=model, optimizer=optimizer, global_step=optimizer.iterations) sub_model_checkpoint = tf.train.Checkpoint( model=sub_model, global_step=optimizer.iterations) if sub_model_export_name else None latest_checkpoint_file = tf.train.latest_checkpoint(model_dir) if latest_checkpoint_file: logging.info('Checkpoint file %s found and restoring from ' 'checkpoint', latest_checkpoint_file) checkpoint.restore(latest_checkpoint_file) logging.info('Loading from checkpoint file completed') current_step = optimizer.iterations.numpy() checkpoint_name = 'ctl_step_{step}.ckpt' logs = {} callback_list.on_train_begin() while current_step < total_training_steps and not model.stop_training: if current_step % steps_per_epoch == 0: callback_list.on_epoch_begin(int(current_step / steps_per_epoch) + 1) # Training loss/metric are taking average over steps inside micro # training loop. We reset the their values before each round. train_loss_metric.reset_states() for metric in train_metrics + model.metrics: metric.reset_states() callback_list.on_batch_begin(current_step) # Runs several steps in the host while loop. steps = steps_to_run(current_step, steps_between_evals, steps_per_loop) if tf.config.list_physical_devices('GPU'): # TODO(zongweiz): merge with train_steps once tf.while_loop # GPU performance bugs are fixed. for _ in range(steps): train_single_step(train_iterator) else: # Converts steps to a Tensor to avoid tf.function retracing. train_steps(train_iterator, tf.convert_to_tensor(steps, dtype=tf.int32)) train_loss = _float_metric_value(train_loss_metric) current_step += steps # Updates training logging. training_status = 'Train Step: %d/%d / loss = %s' % ( current_step, total_training_steps, train_loss) if current_step >= last_summary_step + train_summary_interval: summary_writer = train_summary_writer last_summary_step = current_step else: summary_writer = tf.summary.create_noop_writer() with summary_writer.as_default(): if callable(optimizer.learning_rate): tf.summary.scalar( 'learning_rate', optimizer.learning_rate(current_step), step=current_step) tf.summary.scalar(train_loss_metric.name, train_loss, step=current_step) for metric in train_metrics + model.metrics: metric_value = _float_metric_value(metric) training_status += ' %s = %f' % (metric.name, metric_value) tf.summary.scalar(metric.name, metric_value, step=current_step) summary_writer.flush() logging.info(training_status) # If no need for evaluation, we only call on_batch_end with train_loss, # this is to ensure we get granular global_step/sec on Tensorboard. if current_step % steps_between_evals: callback_list.on_batch_end(current_step - 1, {'loss': train_loss}) else: # Save a submodel with the step in the file name after each epoch. if sub_model_export_name: _save_checkpoint( strategy, sub_model_checkpoint, model_dir, '%s_step_%d.ckpt' % (sub_model_export_name, current_step)) # Save model checkpoints and run validation steps after each epoch # (with the exception of the final epoch which is handled after the # training loop). if current_step < total_training_steps: _save_checkpoint(strategy, checkpoint, model_dir, checkpoint_name.format(step=current_step)) if eval_input_fn: # Re-initialize evaluation metric. eval_loss_metric.reset_states() for metric in eval_metrics + model.metrics: metric.reset_states() logging.info('Running evaluation after step: %s.', current_step) logs = _run_evaluation(current_step, _get_input_iterator(eval_input_fn, strategy)) # We add train_loss here rather than call on_batch_end twice to make # sure that no duplicated values are generated. logs['loss'] = train_loss callback_list.on_batch_end(current_step - 1, logs) # Calls on_epoch_end after each real epoch ends to prevent mis-calculation # of training steps. if current_step % steps_per_epoch == 0: callback_list.on_epoch_end(int(current_step / steps_per_epoch), logs) if sub_model_export_name: _save_checkpoint(strategy, sub_model_checkpoint, model_dir, '%s.ckpt' % sub_model_export_name) _save_checkpoint(strategy, checkpoint, model_dir, checkpoint_name.format(step=current_step)) if eval_input_fn: # Re-initialize evaluation metric. eval_loss_metric.reset_states() for metric in eval_metrics + model.metrics: metric.reset_states() logging.info('Running final evaluation after training is complete.') logs = _run_evaluation(current_step, _get_input_iterator(eval_input_fn, strategy)) callback_list.on_epoch_end(int(current_step / steps_per_epoch), logs) training_summary = { 'total_training_steps': total_training_steps, 'train_loss': _float_metric_value(train_loss_metric), } for metric in model.metrics: training_summary[metric.name] = _float_metric_value(metric) if eval_metrics: training_summary['last_train_metrics'] = _float_metric_value( train_metrics[0]) training_summary['eval_metrics'] = _float_metric_value(eval_metrics[0]) write_txt_summary(training_summary, summary_dir) if not _should_export_summary(strategy): tf.io.gfile.rmtree(summary_dir) callback_list.on_train_end() return model
def run(flags_obj): """Run ResNet ImageNet training and eval loop using custom training loops. Args: flags_obj: An object containing parsed flag values. Raises: ValueError: If fp16 is passed as it is not currently supported. Returns: Dictionary of training and eval stats. """ keras_utils.set_session_config() performance.set_mixed_precision_policy(flags_core.get_tf_dtype(flags_obj)) if tf.config.list_physical_devices('GPU'): if flags_obj.tf_gpu_thread_mode: keras_utils.set_gpu_thread_mode_and_count( per_gpu_thread_count=flags_obj.per_gpu_thread_count, gpu_thread_mode=flags_obj.tf_gpu_thread_mode, num_gpus=flags_obj.num_gpus, datasets_num_private_threads=flags_obj. datasets_num_private_threads) common.set_cudnn_batchnorm_mode() data_format = flags_obj.data_format if data_format is None: data_format = ('channels_first' if tf.config.list_physical_devices('GPU') else 'channels_last') tf.keras.backend.set_image_data_format(data_format) strategy = distribute_utils.get_distribution_strategy( distribution_strategy=flags_obj.distribution_strategy, num_gpus=flags_obj.num_gpus, all_reduce_alg=flags_obj.all_reduce_alg, num_packs=flags_obj.num_packs, tpu_address=flags_obj.tpu) per_epoch_steps, train_epochs, eval_steps = get_num_train_iterations( flags_obj) if flags_obj.steps_per_loop is None: steps_per_loop = per_epoch_steps elif flags_obj.steps_per_loop > per_epoch_steps: steps_per_loop = per_epoch_steps logging.warn('Setting steps_per_loop to %d to respect epoch boundary.', steps_per_loop) else: steps_per_loop = flags_obj.steps_per_loop logging.info( 'Training %d epochs, each epoch has %d steps, ' 'total steps: %d; Eval %d steps', train_epochs, per_epoch_steps, train_epochs * per_epoch_steps, eval_steps) time_callback = keras_utils.TimeHistory( flags_obj.batch_size, flags_obj.log_steps, logdir=flags_obj.model_dir if flags_obj.enable_tensorboard else None) with distribute_utils.get_strategy_scope(strategy): runnable = resnet_runnable.ResnetRunnable(flags_obj, time_callback, per_epoch_steps) eval_interval = flags_obj.epochs_between_evals * per_epoch_steps checkpoint_interval = (steps_per_loop * 5 if flags_obj.enable_checkpoint_and_export else None) summary_interval = steps_per_loop if flags_obj.enable_tensorboard else None checkpoint_manager = tf.train.CheckpointManager( runnable.checkpoint, directory=flags_obj.model_dir, max_to_keep=10, step_counter=runnable.global_step, checkpoint_interval=checkpoint_interval) resnet_controller = orbit.Controller( strategy=strategy, trainer=runnable, evaluator=runnable if not flags_obj.skip_eval else None, global_step=runnable.global_step, steps_per_loop=steps_per_loop, checkpoint_manager=checkpoint_manager, summary_interval=summary_interval, summary_dir=flags_obj.model_dir, eval_summary_dir=os.path.join(flags_obj.model_dir, 'eval')) time_callback.on_train_begin() if not flags_obj.skip_eval: resnet_controller.train_and_evaluate(train_steps=per_epoch_steps * train_epochs, eval_steps=eval_steps, eval_interval=eval_interval) else: resnet_controller.train(steps=per_epoch_steps * train_epochs) time_callback.on_train_end() stats = build_stats(runnable, time_callback) return stats
def run(flags_obj, datasets_override=None, strategy_override=None): """Run MNIST model training and eval loop using native Keras APIs. Args: flags_obj: An object containing parsed flag values. datasets_override: A pair of `tf.data.Dataset` objects to train the model, representing the train and test sets. strategy_override: A `tf.distribute.Strategy` object to use for model. Returns: Dictionary of training and eval stats. """ # Start TF profiler server. tf.profiler.experimental.server.start(flags_obj.profiler_port) strategy = strategy_override or distribute_utils.get_distribution_strategy( distribution_strategy=flags_obj.distribution_strategy, num_gpus=flags_obj.num_gpus, tpu_address=flags_obj.tpu) strategy_scope = distribute_utils.get_strategy_scope(strategy) mnist = tfds.builder('mnist', data_dir=flags_obj.data_dir) if flags_obj.download: mnist.download_and_prepare() mnist_train, mnist_test = datasets_override or mnist.as_dataset( split=['train', 'test'], decoders={'image': decode_image()}, # pylint: disable=no-value-for-parameter as_supervised=True) train_input_dataset = mnist_train.cache().repeat().shuffle( buffer_size=50000).batch(flags_obj.batch_size) eval_input_dataset = mnist_test.cache().repeat().batch( flags_obj.batch_size) with strategy_scope: lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay( 0.05, decay_steps=100000, decay_rate=0.96) optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule) model = build_model() model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['sparse_categorical_accuracy']) num_train_examples = mnist.info.splits['train'].num_examples train_steps = num_train_examples // flags_obj.batch_size train_epochs = flags_obj.train_epochs ckpt_full_path = os.path.join(flags_obj.model_dir, 'model.ckpt-{epoch:04d}') callbacks = [ tf.keras.callbacks.ModelCheckpoint(ckpt_full_path, save_weights_only=True), tf.keras.callbacks.TensorBoard(log_dir=flags_obj.model_dir), ] num_eval_examples = mnist.info.splits['test'].num_examples num_eval_steps = num_eval_examples // flags_obj.batch_size history = model.fit(train_input_dataset, epochs=train_epochs, steps_per_epoch=train_steps, callbacks=callbacks, validation_steps=num_eval_steps, validation_data=eval_input_dataset, validation_freq=flags_obj.epochs_between_evals) export_path = os.path.join(flags_obj.model_dir, 'saved_model') model.save(export_path, include_optimizer=False) eval_output = model.evaluate(eval_input_dataset, steps=num_eval_steps, verbose=2) stats = common.build_stats(history, eval_output, callbacks) return stats