def main(enable_v2_behavior=True):
"""All-in-one main function for tf.distribute tests."""
if enable_v2_behavior:
v2_compat.enable_v2_behavior()
else:
v2_compat.disable_v2_behavior()
# TODO(b/131360402): configure default logical devices.
multi_process_runner.test_main()
def main(enable_v2_behavior=True, config_logical_devices=True):
"""All-in-one main function for tf.distribute tests."""
if config_logical_devices:
app.call_after_init(_set_logical_devices)
if enable_v2_behavior:
v2_compat.enable_v2_behavior()
else:
v2_compat.disable_v2_behavior()
multi_process_runner.test_main()
"and", "fire"
]]
vocab_dataset = dataset_ops.Dataset.from_tensors(vocab_data)
input_array = np.array([["earth", "wind", "and", "fire"],
["fire", "and", "earth", "michigan"]])
input_dataset = dataset_ops.Dataset.from_tensor_slices(input_array).batch(
2, drop_remainder=True)
expected_output = [[3, 4, 5, 6], [6, 5, 3, 1]]
config.set_soft_device_placement(True)
with strategy.scope():
input_data = keras.Input(shape=(None,), dtype=dtypes.string)
layer = index_lookup.IndexLookup(
max_tokens=None,
num_oov_indices=2,
mask_token="",
oov_token="[OOV]",
dtype=dtypes.string)
layer.adapt(vocab_dataset)
int_data = layer(input_data)
model = keras.Model(inputs=input_data, outputs=int_data)
model.compile(loss="mse")
output_dataset = model.predict(input_dataset)
self.assertAllEqual(expected_output, output_dataset)
if __name__ == "__main__":
v2_compat.enable_v2_behavior()
multi_process_runner.test_main()
def main(): """Tests must call this main().""" return multi_process_runner.test_main()
class EpochCounterCallback(callbacks.Callback):
def on_epoch_begin(self, epoch, logs):
self.last_epoch = epoch
model, _, train_ds, steps = _model_setup(test_obj, file_format='')
epoch_counter_cbk = EpochCounterCallback()
cbks = [
callbacks.EarlyStopping(monitor='loss',
min_delta=0.05,
patience=1,
verbose=1), epoch_counter_cbk
]
# Empirically, it is expected that `model.fit()` terminates around the
# 22th epoch. Asserting that it should have been stopped before the 50th
# epoch to avoid flakiness and be more predictable.
model.fit(x=train_ds,
epochs=100,
steps_per_epoch=steps,
callbacks=cbks)
test_obj.assertLess(epoch_counter_cbk.last_epoch, 50)
multi_process_runner.run(
proc_early_stopping,
cluster_spec=test_base.create_cluster_spec(num_workers=2),
args=(self, ))
if __name__ == '__main__':
multi_process_runner.test_main(barrier_parties=2)
multi_process_runner.barrier().wait()
tensor_shape = [2] * num_dims
def variable_fn():
with ops.device(self._local_device):
# The initial value will be broadcasted from worker 0 to others.
initial_value = (array_ops.ones(tensor_shape) if worker_id == 0 else
array_ops.zeros(tensor_shape))
var = variable_scope.get_variable(name='x', initializer=initial_value)
return array_ops.identity(var)
t_out = strategy.extended.call_for_each_replica(variable_fn)
expected_out = np.ones(tensor_shape)
self.assertAllClose(t_out, expected_out)
def worker_fn():
self._maybe_setup_gpus()
tf_config = json.loads(os.environ['TF_CONFIG'])
worker_id = tf_config['task']['index']
for i in range(20):
worker_step_fn(worker_id, num_dims=(i + 1))
with test_util.skip_if_error(self, errors_impl.UnavailableError):
multi_process_runner.run(
worker_fn,
cluster_spec=test_base.create_cluster_spec(num_workers=NUM_WORKERS))
if __name__ == '__main__':
multi_process_runner.test_main(barrier_parties=NUM_WORKERS)
