(episodes_dir,
                               train_suffix=utils.TRAIN_SUFFIX,
                               test_suffix=utils.TEST_SUFFIX)
return dataset_of_datasets, len(test_files)
This function reads all episodes from a given directory.

  Additionally it returns total number of episodes available in one epoch.
  Args:
    episodes_dir: str, directory that has tf_record files.
    train_suffix: str, used during dumping of episodes to indicate training
      records. Default value should be kept unless the default is overwritten
      during creation of these episodes.
    test_suffix: str, used during dumping of episodes to indicate test records.
      Default value should be kept unless the default is overwritten during
      creation of these episodes.

  Returns:
    tf.data.Dataset, that returns a tuple of training and test datasets. Each
      dataset has `image` and `label` fields.
    int, number of episodes read.
  Raises:
    RuntimeError: when some episodes are missing.
This function reads all episodes from a given directory.
This function reads all episodes from a given directory .
def read_episodes_from_records(episodes_dir,
                               train_suffix=utils.TRAIN_SUFFIX,
                               test_suffix=utils.TEST_SUFFIX):
  """This function reads all episodes from a given directory.

  Additionally it returns total number of episodes available in one epoch.
  Args:
    episodes_dir: str, directory that has tf_record files.
    train_suffix: str, used during dumping of episodes to indicate training
      records. Default value should be kept unless the default is overwritten
      during creation of these episodes.
    test_suffix: str, used during dumping of episodes to indicate test records.
      Default value should be kept unless the default is overwritten during
      creation of these episodes.

  Returns:
    tf.data.Dataset, that returns a tuple of training and test datasets. Each
      dataset has `image` and `label` fields.
    int, number of episodes read.
  Raises:
    RuntimeError: when some episodes are missing.
  """
  all_files = sorted(tf.io.gfile.listdir(episodes_dir))
  train_files = [f for f in all_files if f.endswith(train_suffix)]
  test_files = [f for f in all_files if f.endswith(test_suffix)]
  for test_file, train_file in zip(test_files, train_files):
    # Check whether ids match: expected format episode-0001-train.tfrecords
    # TODO(evcu) maybe use regex.
    if test_file.split("-")[1] != train_file.split("-")[1]:
      test_id = int(test_file.split("-")[1])
      train_id = int(train_file.split("-")[1])
      if test_id < train_id:
        raise RuntimeError("Train data missing for %d th episode." % test_id)
      else:
        raise RuntimeError("Test data missing for %d th episode." % train_id)

  # Load episode information to obtain total number of images for each episode.
  info_path = utils.get_info_path(episodes_dir)
  with tf.io.gfile.GFile(info_path, "r") as f:
    all_info = json.load(f)

  def get_total_img_count(file_name):
    return sum(all_info[file_name].values())

  # Define function to load individual tf-record files.
  def _load_and_batch(file_path, n_images):
    dataset = tf.data.TFRecordDataset(file_path)
    dataset = dataset.batch(tf.cast(n_images, tf.int64))
    batch_data = dataset.make_initializable_iterator().get_next()
    # This pipeline doesn't have original label ids.
    return batch_data

  train_paths = tf.data.Dataset.from_tensor_slices(
      [os.path.join(episodes_dir, f) for f in train_files])
  train_n_images = tf.data.Dataset.from_tensor_slices(
      [get_total_img_count(f) for f in train_files])
  test_paths = tf.data.Dataset.from_tensor_slices(
      [os.path.join(episodes_dir, f) for f in test_files])
  test_n_images = tf.data.Dataset.from_tensor_slices(
      [get_total_img_count(f) for f in test_files])
  dataset_of_datasets = tf.data.Dataset.zip(
      (tf.data.Dataset.zip((train_paths, train_n_images)).map(_load_and_batch),
       tf.data.Dataset.zip((test_paths, test_n_images)).map(_load_and_batch)))
  return dataset_of_datasets, len(test_files)
def read_episodes_from_records ( episodes_dir , train_suffix = utils . TRAIN_SUFFIX , test_suffix = utils . TEST_SUFFIX ) : all_files = sorted ( tf . io . gfile . listdir ( episodes_dir ) ) train_files = [ f for f in all_files if f . endswith ( train_suffix ) ] test_files = [ f for f in all_files if f . endswith ( test_suffix ) ] for test_file , train_file in zip ( test_files , train_files ) : # Check whether ids match: expected format episode-0001-train.tfrecords # TODO(evcu) maybe use regex. if test_file . split ( "-" ) [ 1 ] != train_file . split ( "-" ) [ 1 ] : test_id = int ( test_file . split ( "-" ) [ 1 ] ) train_id = int ( train_file . split ( "-" ) [ 1 ] ) if test_id < train_id : raise RuntimeError ( "Train data missing for %d th episode." % test_id ) else : raise RuntimeError ( "Test data missing for %d th episode." % train_id ) # Load episode information to obtain total number of images for each episode. info_path = utils . get_info_path ( episodes_dir ) with tf . io . gfile . GFile ( info_path , "r" ) as f : all_info = json . load ( f ) def get_total_img_count ( file_name ) : return sum ( all_info [ file_name ] . values ( ) ) # Define function to load individual tf-record files. def _load_and_batch ( file_path , n_images ) : dataset = tf . data . TFRecordDataset ( file_path ) dataset = dataset . batch ( tf . cast ( n_images , tf . int64 ) ) batch_data = dataset . make_initializable_iterator ( ) . get_next ( ) # This pipeline doesn't have original label ids. return batch_data train_paths = tf . data . Dataset . from_tensor_slices ( [ os . path . join ( episodes_dir , f ) for f in train_files ] ) train_n_images = tf . data . Dataset . from_tensor_slices ( [ get_total_img_count ( f ) for f in train_files ] ) test_paths = tf . data . Dataset . from_tensor_slices ( [ os . path . join ( episodes_dir , f ) for f in test_files ] ) test_n_images = tf . data . Dataset . from_tensor_slices ( [ get_total_img_count ( f ) for f in test_files ] ) dataset_of_datasets = tf . data . Dataset . zip ( ( tf . data . Dataset . zip ( ( train_paths , train_n_images ) ) . map ( _load_and_batch ) , tf . data . Dataset . zip ( ( test_paths , test_n_images ) ) . map ( _load_and_batch ) ) ) return dataset_of_datasets , len ( test_files )