def get_eval_data_iters(dataset: dataset_utils.Dataset,
                        config: ml_collections.ConfigDict,
                        is_final_eval: bool = False):
  """Get data iterators with number of steps for evaluation."""
  data_iters = collections.OrderedDict()
  batch_size_multiplier = 1
  for data_split in config.eval.data_splits.split(','):
    # Get iterator for the split with total examples.
    if data_split == 'train':
      num_examples_per_epoch = dataset.meta_data['num_train_examples']
      data_iter = dataset.train_iter
    elif data_split == 'validation':
      num_examples_per_epoch = dataset.meta_data['num_eval_examples']
      if config.dataset_configs.get('num_eval_clips', 1) > 1:
        batch_size_multiplier *= config.dataset_configs['num_eval_clips']
      data_iter = dataset.valid_iter
    elif data_split == 'test':
      num_examples_per_epoch = dataset.meta_data['num_test_examples']
      if config.dataset_configs.get('num_test_clips', 1) > 1:
        batch_size_multiplier *= config.dataset_configs['num_test_clips']
      if config.dataset_configs.get('do_three_spatial_crops', False):
        batch_size_multiplier *= 3
      data_iter = dataset.test_iter
    else:
      raise ValueError(f'Invalid data split: {data_split}')

    if config.eval.get('num_examples') is not None:
      total_num_examples = config.eval.num_examples
    else:
      total_num_examples = num_examples_per_epoch

    num_repeats = 1
    # Run more at final-step evaluation
    if is_final_eval:
      num_repeats = config.eval.get('final_num_repeats', 1)
      num_example_multiplier = config.eval.get('final_num_example_multiplier',
                                               1)
      total_num_examples *= num_repeats * num_example_multiplier

    # Add a few steps to account for the partial batch at epoch end
    batch_size = batch_size_multiplier * config.eval_batch_size
    if num_examples_per_epoch > 0:
      valid_fraction = num_examples_per_epoch / (
          math.ceil(num_examples_per_epoch / batch_size) * batch_size)
    else:
      valid_fraction = 1.
    num_steps = int(
        math.ceil(total_num_examples / config.eval_batch_size / valid_fraction))
    data_iters[data_split] = (data_iter, num_steps, total_num_examples,
                              num_repeats)
  return data_iters
Get data iterators with number of steps for evaluation.
videogvt/train_lib/eval_utils.py