(dataset_spec, split, pool, sampler)
Generates a stream of compact episode representations.

  Each episode is chunked into:

  * a "flush" chunk, which is meant to allow to flush examples, in case we are
    at the end of an epoch for one or more class in the episode (we want to
    avoid accidentally repeating an example due to epoch boundaries), and
  * some number of additional chunks (for example, a "support" chunk and a
    "query" chunk).

  To make sure the input pipeline knows where the episode boundary is within the
  stream (and where the boundary is between chunks in an episode), we enforce
  that each chunk has a fixed size by padding with placeholder dataset IDs (of
  value `num_classes`) as needed (in some cases it's possible that no padding is
  ever needed). The size of each chunk is prescribed by the
  `compute_chunk_sizes` method of `sampler`, which also implicitly defines the
  number of additional chunks (i.e. `len(chunk_sizes) - 1`).

  Instead of explicitly representing all elements of the dataset ID stream, this
  generator returns a compact representation where repeated elements are
  replaced with a `(dataset_id, num_repeats)` tuple.

  This generator is meant to be used with
  `tf.data.experimental.choose_from_datasets` and assumes that the list of
  tf.data.Dataset objects corresponding to each class in the dataset (there are
  `num_classes` of them, which is determined by inspecting the `dataset_spec`
  argument using the `split` argument) is appended with a placeholder Dataset
  (which has index `num_classes` in the list) which outputs a constant `(b'',
  PLACEHOLDER_CLASS_ID)` tuple).

  Note that a dataset ID is different from the (absolute) class ID: the dataset
  ID refers to the index of the Dataset in the list of Dataset objects, and the
  class ID (or label) refers to the second element of the tuple that the Dataset
  outputs.

  Args:
    dataset_spec: DatasetSpecification, dataset specification.
    split: one of Split.TRAIN, Split.VALID, or Split.TEST.
    pool: A string ('train' or 'test') or None, indicating which example-level
      split to select, if the current dataset has them.
    sampler: EpisodeDescriptionSampler instance.

  Yields:
    episode_representation: tensor of shape [N, 2], where N varies dynamically
      between episodes.
Generates a stream of compact episode representations.
Generates a stream of compact episode representations .
def episode_representation_generator(dataset_spec, split, pool, sampler):
  """Generates a stream of compact episode representations.

  Each episode is chunked into:

  * a "flush" chunk, which is meant to allow to flush examples, in case we are
    at the end of an epoch for one or more class in the episode (we want to
    avoid accidentally repeating an example due to epoch boundaries), and
  * some number of additional chunks (for example, a "support" chunk and a
    "query" chunk).

  To make sure the input pipeline knows where the episode boundary is within the
  stream (and where the boundary is between chunks in an episode), we enforce
  that each chunk has a fixed size by padding with placeholder dataset IDs (of
  value `num_classes`) as needed (in some cases it's possible that no padding is
  ever needed). The size of each chunk is prescribed by the
  `compute_chunk_sizes` method of `sampler`, which also implicitly defines the
  number of additional chunks (i.e. `len(chunk_sizes) - 1`).

  Instead of explicitly representing all elements of the dataset ID stream, this
  generator returns a compact representation where repeated elements are
  replaced with a `(dataset_id, num_repeats)` tuple.

  This generator is meant to be used with
  `tf.data.experimental.choose_from_datasets` and assumes that the list of
  tf.data.Dataset objects corresponding to each class in the dataset (there are
  `num_classes` of them, which is determined by inspecting the `dataset_spec`
  argument using the `split` argument) is appended with a placeholder Dataset
  (which has index `num_classes` in the list) which outputs a constant `(b'',
  PLACEHOLDER_CLASS_ID)` tuple).

  Note that a dataset ID is different from the (absolute) class ID: the dataset
  ID refers to the index of the Dataset in the list of Dataset objects, and the
  class ID (or label) refers to the second element of the tuple that the Dataset
  outputs.

  Args:
    dataset_spec: DatasetSpecification, dataset specification.
    split: one of Split.TRAIN, Split.VALID, or Split.TEST.
    pool: A string ('train' or 'test') or None, indicating which example-level
      split to select, if the current dataset has them.
    sampler: EpisodeDescriptionSampler instance.

  Yields:
    episode_representation: tensor of shape [N, 2], where N varies dynamically
      between episodes.
  """
  chunk_sizes = sampler.compute_chunk_sizes()
  # An episode always starts with a "flush" chunk to allow flushing examples at
  # class epoch boundaries, and contains `len(chunk_sizes) - 1` additional
  # chunks.
  flush_chunk_size, other_chunk_sizes = chunk_sizes[0], chunk_sizes[1:]

  class_set = dataset_spec.get_classes(split)
  num_classes = len(class_set)
  placeholder_dataset_id = num_classes

  total_images_per_class = dict(
      (class_idx,
       dataset_spec.get_total_images_per_class(class_set[class_idx], pool))
      for class_idx in range(num_classes))
  cursors = [0] * num_classes

  run_counter = 0
  # Infinite loop over episodes.
  while True:
    flushed_dataset_indices = []
    selected_dataset_indices = [[] for _ in other_chunk_sizes]
    if run_counter == 0:
      # Sample an episode description. A description is a tuple of
      # `(class_idx, ...)` tuples, where `class_idx` indicates the class to
      # sample from and the remaining `len(chunk_sizes) - 1` elements indicate
      # how many examples to allocate to each chunk.
      episode_description = sampler.sample_episode_description()

    if run_counter < sampler.episode_description_switch_frequency - 1:
      run_counter += 1
    else:
      run_counter = 0

    for element in episode_description:
      class_idx, distribution = element[0], element[1:]
      total_requested = sum(distribution)
      if total_requested > total_images_per_class[class_idx]:
        raise ValueError("Requesting more images than what's available for the "
                         'whole class')
      # If the total number of requested examples is greater than the number of
      # examples remaining for the current pass over class `class_idx`, we flush
      # the remaining examples and start a new pass over class `class_idx`.
      # TODO(lamblinp): factor this out into its own tracker class for
      # readability and testability.
      remaining = total_images_per_class[class_idx] - cursors[class_idx]
      if total_requested > remaining:
        flushed_dataset_indices.append([class_idx, remaining])
        cursors[class_idx] = 0
      # Elements of `distribution` correspond to how many examples of class
      # `class_idx` to allocate for each chunk (e.g. in a few-shot learning
      # context `distribution = [5, 8]` would allocate 5 examples to the
      # "support" chunk and 8 examples to the "query" chunk). Elements of
      # `selected_dataset_indices` correspond to the list of dataset indices
      # that have so far been requested for each chunk.
      for num_to_allocate, dataset_indices in zip(distribution,
                                                  selected_dataset_indices):
        dataset_indices.append([class_idx, num_to_allocate])
      cursors[class_idx] += total_requested

    # An episode sequence is generated in multiple phases, each padded with an
    # agreed-upon number of placeholder dataset IDs.

    _pad(flushed_dataset_indices, flush_chunk_size, placeholder_dataset_id)
    for dataset_indices, chunk_size in zip(selected_dataset_indices,
                                           other_chunk_sizes):
      _pad(dataset_indices, chunk_size, placeholder_dataset_id)

    episode_representation = np.array(
        list(
            itertools.chain(flushed_dataset_indices,
                            *selected_dataset_indices)),
        dtype='int64')
    yield episode_representation
def episode_representation_generator ( dataset_spec , split , pool , sampler ) : chunk_sizes = sampler . compute_chunk_sizes ( ) # An episode always starts with a "flush" chunk to allow flushing examples at # class epoch boundaries, and contains `len(chunk_sizes) - 1` additional # chunks. flush_chunk_size , other_chunk_sizes = chunk_sizes [ 0 ] , chunk_sizes [ 1 : ] class_set = dataset_spec . get_classes ( split ) num_classes = len ( class_set ) placeholder_dataset_id = num_classes total_images_per_class = dict ( ( class_idx , dataset_spec . get_total_images_per_class ( class_set [ class_idx ] , pool ) ) for class_idx in range ( num_classes ) ) cursors = [ 0 ] * num_classes run_counter = 0 # Infinite loop over episodes. while True : flushed_dataset_indices = [ ] selected_dataset_indices = [ [ ] for _ in other_chunk_sizes ] if run_counter == 0 : # Sample an episode description. A description is a tuple of # `(class_idx, ...)` tuples, where `class_idx` indicates the class to # sample from and the remaining `len(chunk_sizes) - 1` elements indicate # how many examples to allocate to each chunk. episode_description = sampler . sample_episode_description ( ) if run_counter < sampler . episode_description_switch_frequency - 1 : run_counter += 1 else : run_counter = 0 for element in episode_description : class_idx , distribution = element [ 0 ] , element [ 1 : ] total_requested = sum ( distribution ) if total_requested > total_images_per_class [ class_idx ] : raise ValueError ( "Requesting more images than what's available for the " 'whole class' ) # If the total number of requested examples is greater than the number of # examples remaining for the current pass over class `class_idx`, we flush # the remaining examples and start a new pass over class `class_idx`. # TODO(lamblinp): factor this out into its own tracker class for # readability and testability. remaining = total_images_per_class [ class_idx ] - cursors [ class_idx ] if total_requested > remaining : flushed_dataset_indices . append ( [ class_idx , remaining ] ) cursors [ class_idx ] = 0 # Elements of `distribution` correspond to how many examples of class # `class_idx` to allocate for each chunk (e.g. in a few-shot learning # context `distribution = [5, 8]` would allocate 5 examples to the # "support" chunk and 8 examples to the "query" chunk). Elements of # `selected_dataset_indices` correspond to the list of dataset indices # that have so far been requested for each chunk. for num_to_allocate , dataset_indices in zip ( distribution , selected_dataset_indices ) : dataset_indices . append ( [ class_idx , num_to_allocate ] ) cursors [ class_idx ] += total_requested # An episode sequence is generated in multiple phases, each padded with an # agreed-upon number of placeholder dataset IDs. _pad ( flushed_dataset_indices , flush_chunk_size , placeholder_dataset_id ) for dataset_indices , chunk_size in zip ( selected_dataset_indices , other_chunk_sizes ) : _pad ( dataset_indices , chunk_size , placeholder_dataset_id ) episode_representation = np . array ( list ( itertools . chain ( flushed_dataset_indices , * selected_dataset_indices ) ) , dtype = 'int64' ) yield episode_representation