(
    builder,
    md_version,
    meta_split,
    source_id = None,
    shuffle_seed = None,
    **as_dataset_kwargs)
return dataset
Creates an episode dataset.

  This function creates an episode dataset for a single source. For multi-source
  pipelines, use the `meta_dataset` function.

  Args:
    builder: Meta-Dataset builder.
    md_version: Meta-Dataset md_version in {'v1', 'v2'}.
    meta_split: meta-split (case-insensitive) in {'train', 'valid', 'test'}.
    source_id: source ID to output alongside episodes, if not None.
    shuffle_seed: class dataset shuffle buffer seed.
    **as_dataset_kwargs: kwargs passed to the `as_dataset` method.

  Returns:
    The episode dataset.
Creates an episode dataset.
Creates an episode dataset .
def episode_dataset(
    builder,
    md_version,
    meta_split,
    source_id = None,
    shuffle_seed = None,
    **as_dataset_kwargs):
  """Creates an episode dataset.

  This function creates an episode dataset for a single source. For multi-source
  pipelines, use the `meta_dataset` function.

  Args:
    builder: Meta-Dataset builder.
    md_version: Meta-Dataset md_version in {'v1', 'v2'}.
    meta_split: meta-split (case-insensitive) in {'train', 'valid', 'test'}.
    source_id: source ID to output alongside episodes, if not None.
    shuffle_seed: class dataset shuffle buffer seed.
    **as_dataset_kwargs: kwargs passed to the `as_dataset` method.

  Returns:
    The episode dataset.
  """

  dataset_spec = dataset_spec_lib.as_dataset_spec(
      builder.info.metadata['dataset_specs'][md_version])
  data_config = config_lib.DataConfig()
  episode_description_config = config_lib.EpisodeDescriptionConfig()

  def _as_class_dataset(args):
    relative_label, seed = args
    seed = None if seed is None else seed[0]
    class_dataset = builder.as_class_dataset(
        md_version=md_version,
        meta_split=meta_split,
        relative_label=relative_label,
        shuffle_buffer_size=data_config.shuffle_buffer_size,
        shuffle_seed=seed,
        num_prefetch=data_config.num_prefetch,
        decoders={'image': tfds.decode.SkipDecoding()},
        read_config=tfds.ReadConfig(try_autocache=False),
        as_supervised=True,
        **as_dataset_kwargs
    ).repeat()
    logging.info("Created class %d for %s's meta-%s split",
                 relative_label, builder.builder_config.name, meta_split)
    return class_dataset

  num_classes = len(dataset_spec.get_classes(getattr(learning_spec.Split,
                                                     meta_split.upper())))
  # If a shuffle seed is passed, we split it into `num_classes` independent
  # shufle seeds so that each class datasets' shuffle buffer is seeded
  # differently. Sharing random seeds across shuffle buffers is considered bad
  # practice because it can introduce correlations across random sequences of
  # examples for different classes.
  shuffle_seeds = ([None] * num_classes if shuffle_seed is None else
                   tfre.stateless_split(create_rng_state(shuffle_seed),
                                        num_classes))
  with concurrent.futures.ThreadPoolExecutor(
      max_workers=_NUM_WORKERS) as executor:
    class_datasets = list(executor.map(_as_class_dataset,
                                       enumerate(shuffle_seeds)))

  placeholder_id = reader.PLACEHOLDER_CLASS_ID
  class_datasets.append(tf.data.Dataset.zip((
      tf.data.Dataset.from_tensors(b'').repeat(),
      tf.data.Dataset.from_tensors(tf.cast(placeholder_id, tf.int64)).repeat()
  )))

  sampler = sampling.EpisodeDescriptionSampler(
      dataset_spec,
      getattr(learning_spec.Split, meta_split.upper()),
      episode_description_config,
      use_dag_hierarchy=builder.builder_config.name == 'ilsvrc_2012',
      use_bilevel_hierarchy=builder.builder_config.name == 'omniglot')
  chunk_sizes = sampler.compute_chunk_sizes()

  dataset = tf.data.Dataset.choose_from_datasets(
      class_datasets,
      tf.data.Dataset.from_generator(
          functools.partial(
              reader.episode_representation_generator,
              dataset_spec=dataset_spec,
              split=getattr(learning_spec.Split, meta_split.upper()),
              pool=None,
              sampler=sampler),
          tf.int64,
          tf.TensorShape([None, 2]),
      ).map(reader.decompress_episode_representation).unbatch()
  ).batch(
      sum(chunk_sizes)
  ).prefetch(
      1
  ).map(
      functools.partial(
          pipeline.process_episode,
          chunk_sizes=chunk_sizes,
          image_size=data_config.image_height,
          simclr_episode_fraction=(
              episode_description_config.simclr_episode_fraction))
  )

  if source_id is not None:
    dataset = tf.data.Dataset.zip(
        (dataset, tf.data.Dataset.from_tensors(source_id).repeat()))

  return dataset
def episode_dataset ( builder , md_version , meta_split , source_id = None , shuffle_seed = None , * * as_dataset_kwargs ) : dataset_spec = dataset_spec_lib . as_dataset_spec ( builder . info . metadata [ 'dataset_specs' ] [ md_version ] ) data_config = config_lib . DataConfig ( ) episode_description_config = config_lib . EpisodeDescriptionConfig ( ) def _as_class_dataset ( args ) : relative_label , seed = args seed = None if seed is None else seed [ 0 ] class_dataset = builder . as_class_dataset ( md_version = md_version , meta_split = meta_split , relative_label = relative_label , shuffle_buffer_size = data_config . shuffle_buffer_size , shuffle_seed = seed , num_prefetch = data_config . num_prefetch , decoders = { 'image' : tfds . decode . SkipDecoding ( ) } , read_config = tfds . ReadConfig ( try_autocache = False ) , as_supervised = True , * * as_dataset_kwargs ) . repeat ( ) logging . info ( "Created class %d for %s's meta-%s split" , relative_label , builder . builder_config . name , meta_split ) return class_dataset num_classes = len ( dataset_spec . get_classes ( getattr ( learning_spec . Split , meta_split . upper ( ) ) ) ) # If a shuffle seed is passed, we split it into `num_classes` independent # shufle seeds so that each class datasets' shuffle buffer is seeded # differently. Sharing random seeds across shuffle buffers is considered bad # practice because it can introduce correlations across random sequences of # examples for different classes. shuffle_seeds = ( [ None ] * num_classes if shuffle_seed is None else tfre . stateless_split ( create_rng_state ( shuffle_seed ) , num_classes ) ) with concurrent . futures . ThreadPoolExecutor ( max_workers = _NUM_WORKERS ) as executor : class_datasets = list ( executor . map ( _as_class_dataset , enumerate ( shuffle_seeds ) ) ) placeholder_id = reader . PLACEHOLDER_CLASS_ID class_datasets . append ( tf . data . Dataset . zip ( ( tf . data . Dataset . from_tensors ( b'' ) . repeat ( ) , tf . data . Dataset . from_tensors ( tf . cast ( placeholder_id , tf . int64 ) ) . repeat ( ) ) ) ) sampler = sampling . EpisodeDescriptionSampler ( dataset_spec , getattr ( learning_spec . Split , meta_split . upper ( ) ) , episode_description_config , use_dag_hierarchy = builder . builder_config . name == 'ilsvrc_2012' , use_bilevel_hierarchy = builder . builder_config . name == 'omniglot' ) chunk_sizes = sampler . compute_chunk_sizes ( ) dataset = tf . data . Dataset . choose_from_datasets ( class_datasets , tf . data . Dataset . from_generator ( functools . partial ( reader . episode_representation_generator , dataset_spec = dataset_spec , split = getattr ( learning_spec . Split , meta_split . upper ( ) ) , pool = None , sampler = sampler ) , tf . int64 , tf . TensorShape ( [ None , 2 ] ) , ) . map ( reader . decompress_episode_representation ) . unbatch ( ) ) . batch ( sum ( chunk_sizes ) ) . prefetch ( 1 ) . map ( functools . partial ( pipeline . process_episode , chunk_sizes = chunk_sizes , image_size = data_config . image_height , simclr_episode_fraction = ( episode_description_config . simclr_episode_fraction ) ) ) if source_id is not None : dataset = tf . data . Dataset . zip ( ( dataset , tf . data . Dataset . from_tensors ( source_id ) . repeat ( ) ) ) return dataset