def testProperties(self):
input_context = distribute_lib.InputContext(
num_input_pipelines=2, input_pipeline_id=1, num_replicas_in_sync=6)
self.assertEqual(6, input_context.num_replicas_in_sync)
self.assertEqual(1, input_context.input_pipeline_id)
self.assertEqual(2, input_context.num_input_pipelines)
def testPerReplicaBatchSize(self):
input_context = distribute_lib.InputContext(
num_input_pipelines=2, input_pipeline_id=1, num_replicas_in_sync=6)
self.assertEqual(2, input_context.get_per_replica_batch_size(12))
with self.assertRaises(ValueError):
input_context.get_per_replica_batch_size(13)
def _experimental_distribute_datasets_from_function(self, dataset_fn):
return input_lib.get_distributed_datasets_from_function(
dataset_fn,
self._input_workers,
[distribute_lib.InputContext()],
self._container_strategy())
def _make_input_context(self):
input_context = distribute_lib.InputContext(
num_input_pipelines=self._num_workers,
input_pipeline_id=self._id_in_cluster,
num_replicas_in_sync=self._num_replicas_in_sync)
return input_context
def _create_dataset_or_input_fn(self, input_type, input_fn):
if input_type == "input_fn":
return input_fn
else:
return input_fn(distribute_lib.InputContext())
def _make_input_fn_iterator(
self,
input_fn,
replication_mode=distribute_lib.InputReplicationMode.PER_WORKER):
return input_lib.InputFunctionIterator(input_fn, self._input_workers,
[distribute_lib.InputContext()])
def _distribute_datasets_from_function(self, dataset_fn, options):
return input_lib.get_distributed_datasets_from_function(
dataset_fn, self._input_workers_with_options(options),
[distribute_lib.InputContext()], self._container_strategy())
def _distribute_datasets_from_function(self, dataset_fn, options): return dataset_fn(distribute_lib.InputContext())
def _experimental_distribute_datasets_from_function(self, dataset_fn): return dataset_fn(distribute_lib.InputContext())
