Exemplos de load_distributed_dataset em Python

Exemplo n.º 1

Arquivo: variational_inference.py Projeto: zilong-z/edward2

 def clean_test_input_fn(ctx):
     """Sets up local (per-core) dataset batching for testing."""
     dataset = utils.load_distributed_dataset(
         split=tfds.Split.TEST,
         name=FLAGS.dataset,
         batch_size=FLAGS.per_core_batch_size,
         use_bfloat16=FLAGS.use_bfloat16)
     if ctx and ctx.num_input_pipelines > 1:
         dataset = dataset.shard(ctx.num_input_pipelines,
                                 ctx.input_pipeline_id)
     return dataset

Exemplo n.º 2

Arquivo: deterministic.py Projeto: arahosu/edward2

 def train_input_fn(ctx):
   """Sets up local (per-core) dataset batching."""
   dataset = utils.load_distributed_dataset(
       split=tfds.Split.TRAIN,
       name=FLAGS.dataset,
       batch_size=FLAGS.per_core_batch_size,
       drop_remainder=True,
       use_bfloat16=FLAGS.use_bfloat16)
   if ctx and ctx.num_input_pipelines > 1:
     dataset = dataset.shard(ctx.num_input_pipelines, ctx.input_pipeline_id)
   return dataset