def _distribution_strategies():
return [
collective_all_reduce_strategy.CollectiveAllReduceStrategy(),
mirrored_strategy.MirroredStrategy(),
# TODO(pulkitb): Add parameter_server
# parameter_server_strategy.ParameterServerStrategy(),
one_device_strategy.OneDeviceStrategy('/cpu:0'),
]
def testNotPassingASessionInGraph(self):
distribution = one_device_strategy.OneDeviceStrategy("/device:CPU:0")
step_function, _ = single_loss_example(
lambda: gradient_descent.GradientDescentOptimizer(0.2),
distribution)
with context.graph_mode(), ops.Graph().as_default():
with self.assertRaisesRegexp(ValueError, "Should provide"):
_ = monitor_lib.Monitor(step_function, session=None)
def testPassingASessionInEager(self):
distribution = one_device_strategy.OneDeviceStrategy("/device:CPU:0")
step_function, _ = single_loss_example(
lambda: gradient_descent.GradientDescentOptimizer(0.2),
distribution)
with session.Session() as sess, context.eager_mode():
with self.assertRaisesRegexp(ValueError, "Should not provide"):
_ = monitor_lib.Monitor(step_function, sess)
def test_loss_class_as_metric_with_distribution(self):
distribution = one_device_strategy.OneDeviceStrategy('/device:CPU:0')
with distribution.scope():
metric_container = compile_utils.MetricsContainer(
losses_mod.MeanSquaredError())
y_t, y_p = array_ops.ones((10, 5)), array_ops.zeros((10, 5))
metric_container.update_state(y_t, y_p)
self.assertLen(metric_container.metrics, 1)
metric = metric_container.metrics[0]
self.assertEqual(metric.name, 'mean_squared_error')
self.assertEqual(metric.result().numpy(), 1.)
def testMakeInputFnIteratorWithDataset(self):
d = one_device_strategy.OneDeviceStrategy("/device:CPU:0")
dataset_fn = lambda: dataset_ops.Dataset.range(10)
expected_values = [[i] for i in range(10)]
input_fn = self._input_fn_to_test_input_context(
dataset_fn,
expected_num_replicas_in_sync=1,
expected_num_input_pipelines=1,
expected_input_pipeline_id=0)
iterator = d.make_input_fn_iterator(input_fn)
self._test_input_fn_iterator(
iterator, d.extended.worker_devices, expected_values)
def test_cpu_distribution(self):
vocab_data = ["earth", "wind", "and", "fire"]
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
expected_output = [[2, 3, 4, 5], [5, 4, 2, 1]]
strategy = one_device_strategy.OneDeviceStrategy("/cpu:0")
with strategy.scope():
input_data = keras.Input(shape=(None, ), dtype=dtypes.string)
layer = get_layer_class()()
layer.set_vocabulary(vocab_data)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
output_dataset = model.predict(input_array)
self.assertAllEqual(expected_output, output_dataset)
def testMakeInputFnIteratorWithCallable(self):
d = one_device_strategy.OneDeviceStrategy("/device:CPU:0")
def fn():
dataset = dataset_ops.Dataset.range(10)
it = dataset.make_one_shot_iterator()
return it.get_next
expected_values = [[i] for i in range(10)]
input_fn = self._input_fn_to_test_input_context(
fn,
expected_num_replicas_in_sync=1,
expected_num_input_pipelines=1,
expected_input_pipeline_id=0)
iterator = d.make_input_fn_iterator(input_fn)
self._test_input_fn_iterator(
iterator, d.extended.worker_devices, expected_values,
test_reinitialize=False)
return tpu_lib.TPUStrategy(resolver, device_assignment, **kwargs)
else:
return tpu_lib.TPUStrategyV1(resolver, steps_per_run,
device_assignment, **kwargs)
return _create_tpu_strategy
# pylint: disable=g-long-lambda
default_strategy = combinations.NamedDistribution(
"Default",
distribution_strategy_context._get_default_strategy, # pylint: disable=protected-access
required_gpus=None)
one_device_strategy = combinations.NamedDistribution(
"OneDeviceCPU",
lambda: one_device_lib.OneDeviceStrategy("/cpu:0"),
required_gpus=None)
one_device_strategy_gpu = combinations.NamedDistribution(
"OneDeviceGPU",
lambda: one_device_lib.OneDeviceStrategy("/gpu:0"),
required_gpus=1)
one_device_strategy_on_worker_1 = combinations.NamedDistribution(
"OneDeviceOnWorker1CPU",
lambda: one_device_lib.OneDeviceStrategy(
"/job:worker/replica:0/task:1/cpu:0"), # pylint: disable=line-too-long
required_gpus=None)
one_device_strategy_gpu_on_worker_1 = combinations.NamedDistribution(
"OneDeviceOnWorker1GPU",
lambda: one_device_lib.OneDeviceStrategy(
"/job:worker/replica:0/task:1/gpu:0"), # pylint: disable=line-too-long
required_gpus=1)
def _get_distribution_strategy(self):
return one_device_strategy.OneDeviceStrategy("/device:CPU:0")
def _get_distribution_strategy(self):
if "test_gpu" in sys.argv[0]:
return one_device_strategy.OneDeviceStrategy("/device:GPU:0")
else:
return one_device_strategy.OneDeviceStrategy("/device:CPU:0")
strategy = tpu_lib.TPUStrategy(resolver, steps_per_run=steps_per_run,
device_assignment=device_assignment,
**kwargs)
return strategy
return _create_tpu_strategy
# pylint: disable=g-long-lambda
default_strategy = combinations.NamedDistribution(
"Default",
distribution_strategy_context._get_default_strategy, # pylint: disable=protected-access
required_gpus=None)
one_device_strategy = combinations.NamedDistribution(
"OneDeviceCPU",
lambda: one_device_lib.OneDeviceStrategy("/cpu:0"),
required_gpus=None)
one_device_strategy_gpu = combinations.NamedDistribution(
"OneDeviceGPU",
lambda: one_device_lib.OneDeviceStrategy("/gpu:0"),
required_gpus=1)
tpu_strategy = combinations.NamedDistribution(
"TPU", _get_tpu_strategy_creator(steps_per_run=2), required_tpu=True)
tpu_strategy_one_step = combinations.NamedDistribution(
"TPUOneStep", _get_tpu_strategy_creator(steps_per_run=1), required_tpu=True)
tpu_strategy_one_core = combinations.NamedDistribution(
"TPUOneCore",
_get_tpu_strategy_creator(steps_per_run=2, use_single_core=True),
required_tpu=True)
tpu_strategy_one_step_one_core = combinations.NamedDistribution(
"TPUOneStepOneCore",
def create_one_device_strategy():
return one_device_strategy.OneDeviceStrategy('cpu:0')
