def get_num_training_steps(hparams, dataset_metadata):
  """Calculates the total number of training steps and possibly steps_per_epoch.

  The main training loop is based on the number of training steps. Thus, for
  datasets that we want to train based on the number of epochs, we need to
  calculate
  the total number of training steps. This function looks for
  `num_training_steps` in hparams, if it exists it returns that as the total
  step and `None` as `steps_per_epoch`. If num_training_steps doesn't exist,
  then it looks for `num_training_epochs` and given the size of training data
  calculates the total steps and steps_per_epoch. In this computation, we assume
  that drop_remainder=True.

  Args:
    hparams: Hyperparameters.
    dataset_metadata: Meta-data that is generated by the dataset_builder.

  Returns:
    total_steps: int; Total number of training steps.
    steps_per_epoch: int or None; number of steps in every epoch.
  """
  # we either use num_training_epochs or num_training_steps
  if hparams.get('num_training_steps'):
    assert not hparams.get('num_training_epochs')
    return hparams.num_training_steps, None
  else:
    assert hparams.num_training_epochs and not hparams.get('num_training_steps')
    steps_per_epoch = dataset_metadata[
        'num_train_examples'] // hparams.batch_size
    return (steps_per_epoch * hparams.num_training_epochs), steps_per_epoch