def testTrainNetwork(self, distribution, optimizer_fn,
use_callable_loss=True):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
iterator = distribution.distribute_dataset(
dataset_fn).make_one_shot_iterator()
def run_step():
return control_flow_ops.group(distribution.unwrap(
distribution.call_for_each_tower(
model_fn, iterator.get_next(), run_concurrently=layer.built)))
if not context.executing_eagerly():
with self.test_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(distribution.fetch(layer.kernel)))
biases.append(self.evaluate(distribution.fetch(layer.bias)))
error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
def step_fn(ctx, *inputs):
del ctx # Unused
return distribution.group(
distribution.call_for_each_tower(
model_fn, *inputs, run_concurrently=layer.built))
iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
def run_step():
return distribution.run_steps_on_dataset(
step_fn, iterator, iterations=2).run_op
self.evaluate(distribution.initialize())
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(5):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
self.evaluate(distribution.finalize())
error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetworkByCallForEachReplica(self, distribution, optimizer_fn,
use_callable_loss):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
iterator = self._get_iterator(
distribution.distribute_dataset(dataset_fn))
def run_step():
return distribution.group(
distribution.call_for_each_replica(
model_fn, args=(iterator.get_next(), )))
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetworkByCallForEachReplica(self, distribution, optimizer_fn,
use_callable_loss):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
iterator = self._get_iterator(distribution.distribute_dataset(dataset_fn))
def run_step():
return distribution.group(
distribution.call_for_each_replica(
model_fn, args=(iterator.get_next(),)))
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self, distribution, optimizer_fn,
use_callable_loss=True):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
iterator = distribution.make_input_fn_iterator(lambda _: dataset_fn())
def run_step():
return control_flow_ops.group(
distribution.experimental_local_results(
distribution.extended.call_for_each_replica(
model_fn, args=(iterator.get_next(),))))
if not context.executing_eagerly():
with self.cached_session() as sess:
sess.run(iterator.initialize())
run_step = sess.make_callable(run_step())
self.evaluate(variables.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
def step_fn(ctx, inputs):
del ctx # Unused
return distribution.group(
distribution.extended.call_for_each_replica(
model_fn, args=(inputs, )))
iterator = self._get_iterator(distribution, dataset_fn)
def run_step():
return distribution.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=2).run_op
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(5):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
created_variables = []
trainable_variables = []
def appending_creator(next_creator, *args, **kwargs):
v = next_creator(*args, **kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with variable_scope.variable_creator_scope(
appending_creator), distribution.scope():
model_fn, dataset_fn, _ = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=True,
create_optimizer_inside_model_fn=True)
def step_fn(ctx, inputs):
del ctx # Unused
return distribution.group(
distribution.call_for_each_replica(model_fn,
args=(inputs, )))
iterator = self._get_iterator(
distribution.distribute_dataset(dataset_fn))
def run_step():
return distribution.run_steps_on_dataset(step_fn,
iterator,
iterations=1).run_op
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
run_step()
def get_expected_variables(optimizer_fn, num_parameter_devices):
variables_map = {
"GradientDescent": ["dense/kernel", "dense/bias"],
"Adagrad": [
"dense/kernel/Adagrad", "dense/kernel",
"dense/bias/Adagrad", "dense/bias"
]
}
variables = variables_map[optimizer_fn().get_name()]
variables.extend([
v + "/replica_{}".format(replica) for v in variables
for replica in range(1, num_parameter_devices)
])
return set([v + ":0" for v in variables])
self.assertEqual(
get_expected_variables(optimizer_fn,
len(distribution.parameter_devices)),
set(created_variables))
def testTrainNetwork(self,
distribution,
optimizer_fn,
use_callable_loss=True):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
iterator = distribution.make_input_fn_iterator(
lambda _: dataset_fn())
def run_step():
return control_flow_ops.group(
distribution.experimental_local_results(
distribution.extended.call_for_each_replica(
model_fn, args=(iterator.get_next(), ))))
if not context.executing_eagerly():
with self.cached_session() as sess:
sess.run(iterator.initialize())
run_step = sess.make_callable(run_step())
self.evaluate(variables.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
created_variables = []
trainable_variables = []
def appending_creator(next_creator, *args, **kwargs):
v = next_creator(*args, **kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with variable_scope.variable_creator_scope(
appending_creator), distribution.scope():
model_fn, dataset_fn, _ = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=True,
create_optimizer_inside_model_fn=True)
def step_fn(ctx, inputs):
del ctx # Unused
return distribution.group(
distribution.extended.call_for_each_replica(
model_fn, args=(inputs, )))
iterator = self._get_iterator(distribution, dataset_fn)
def run_step():
return distribution.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1).run_op
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
run_step()
def get_expected_variables(optimizer_fn, num_parameter_devices):
optimizer = optimizer_fn()
name = optimizer._name
if isinstance(optimizer, optimizer_v2.OptimizerV2):
variables = VAR_MAP_V2[name]
else:
variables = VAR_MAP_V1[name]
extended_variables = [
v + "/replica_{}".format(replica) for v in variables
for replica in range(1, num_parameter_devices)
]
variables = list(variables) + extended_variables
return set([v + ":0" for v in variables])
self.assertEqual(
get_expected_variables(
optimizer_fn,
len(distribution.extended.parameter_devices)),
set(created_variables))
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
created_variables = []
trainable_variables = []
def appending_creator(next_creator, *args, **kwargs):
v = next_creator(*args, **kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with variable_scope.variable_creator_scope(
appending_creator), distribution.scope():
model_fn, dataset_fn, _ = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=True,
create_optimizer_inside_model_fn=True)
def step_fn(ctx, inputs):
del ctx # Unused
return distribution.group(
distribution.extended.call_for_each_replica(
model_fn, args=(inputs,)))
iterator = self._get_iterator(distribution, dataset_fn)
def run_step():
return distribution.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1).run_op
if not context.executing_eagerly():
with self.cached_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
run_step()
def get_expected_variables(optimizer_fn, num_parameter_devices):
variables_map = {
"GradientDescent": ["dense/kernel", "dense/bias"],
"Adagrad": [
"dense/kernel/Adagrad", "dense/kernel",
"dense/bias/Adagrad", "dense/bias"
]
}
variables = variables_map[optimizer_fn().get_name()]
variables.extend([
v + "/replica_{}".format(replica)
for v in variables
for replica in range(1, num_parameter_devices)
])
return set([v + ":0" for v in variables])
self.assertEqual(
get_expected_variables(optimizer_fn,
len(distribution.extended.parameter_devices)),
set(created_variables))
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
created_variables = []
trainable_variables = []
def appending_creator(next_creator, *args, **kwargs):
v = next_creator(*args, **kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with variable_scope.variable_creator_scope(
appending_creator), distribution.scope():
model_fn, dataset, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=True,
create_optimizer_inside_model_fn=True)
iterator = distribution.distribute_dataset(dataset)
def run_step():
return distribution.group(
distribution.call_for_each_tower(
model_fn,
iterator.get_next(),
run_concurrently=layer.built))
if not context.executing_eagerly():
with self.test_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
run_step()
def get_expected_variables(optimizer_fn, num_parameter_devices):
variables_map = {
"GradientDescent": ["dense/kernel", "dense/bias"],
"Adam": [
"dense/kernel", "dense/bias", "beta1_power",
"beta2_power", "dense/kernel/Adam",
"dense/kernel/Adam_1", "dense/bias/Adam",
"dense/bias/Adam_1"
]
}
variables = variables_map[optimizer_fn().get_name()]
variables.extend([
v + "/replica_{}".format(replica) for v in variables
for replica in range(1, num_parameter_devices)
])
return set([v + ":0" for v in variables])
self.assertEqual(
get_expected_variables(optimizer_fn,
len(distribution.parameter_devices)),
set(created_variables))
def testOptimizerInsideModelFn(self, distribution, optimizer_fn):
created_variables = []
trainable_variables = []
def appending_creator(next_creator, *args, **kwargs):
v = next_creator(*args, **kwargs)
created_variables.append(v.name)
if "trainable" in kwargs and kwargs["trainable"]:
trainable_variables.append(v.name)
return v
# Creator scope needs to be set before it's used inside
# `distribution.scope`.
with variable_scope.variable_creator_scope(
appending_creator), distribution.scope():
model_fn, dataset, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=True,
create_optimizer_inside_model_fn=True)
iterator = distribution.distribute_dataset(dataset)
def run_step():
return distribution.group(
distribution.call_for_each_tower(
model_fn, iterator.get_next(), run_concurrently=layer.built))
if not context.executing_eagerly():
with self.test_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
run_step()
def get_expected_variables(optimizer_fn, num_parameter_devices):
variables_map = {
"GradientDescent": ["dense/kernel", "dense/bias"],
"Adam": [
"dense/kernel", "dense/bias", "beta1_power", "beta2_power",
"dense/kernel/Adam", "dense/kernel/Adam_1", "dense/bias/Adam",
"dense/bias/Adam_1"
]
}
variables = variables_map[optimizer_fn().get_name()]
variables.extend([
v + "/replica_{}".format(replica)
for v in variables
for replica in range(1, num_parameter_devices)
])
return set([v + ":0" for v in variables])
self.assertEqual(
get_expected_variables(optimizer_fn,
len(distribution.parameter_devices)),
set(created_variables))
def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss,
is_tpu):
with distribution.scope():
model_fn, dataset, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
# TODO(isaprykin): Eliminate `is_tpu`. Probably add a
# `DistributionStrategy.create_monitor` so that each DistributionStrategy
# could influence its training loop. That method would return an instance
# of Monitor. TPUMonitor would execute tpu.initialize_system() and
# tpu.shutdown_system().
if is_tpu:
dataset = dataset.batch(2)
iterator = distribution.distribute_dataset(dataset)
def run_step():
# TODO(isaprykin): Make iterator get_next() return a list of sub-
# batches for each iteration. Pass iterator.get_next() and not iterator
# to call_for_each_tower.
return distribution.group(
distribution.call_for_each_tower(
model_fn,
iterator.get_next() if not is_tpu else iterator,
run_concurrently=layer.built))
if not context.executing_eagerly():
with self.test_session() as sess:
if is_tpu:
sess.run(tpu.initialize_system())
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(distribution.fetch(layer.kernel)))
biases.append(self.evaluate(distribution.fetch(layer.bias)))
if is_tpu:
with self.test_session() as sess:
sess.run(tpu.shutdown_system())
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss,
is_tpu):
# TODO(priyag): Remove this once the step TPU Strategy is stable.
if is_tpu:
self.skipTest("TPU tests are WIP.")
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn, use_bias=True, use_callable_loss=use_callable_loss)
# TODO(isaprykin): Eliminate `is_tpu`. Probably add a
# `DistributionStrategy.create_monitor` so that each DistributionStrategy
# could influence its training loop. That method would return an instance
# of Monitor. TPUMonitor would execute tpu.initialize_system() and
# tpu.shutdown_system().
iterator = distribution.distribute_dataset(
dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.group(
distribution.call_for_each_tower(
model_fn, iterator.get_next(), run_concurrently=layer.built))
if not context.executing_eagerly():
with self.test_session() as sess:
if is_tpu:
sess.run(tpu.initialize_system())
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
if is_tpu:
with self.test_session() as sess:
sess.run(tpu.shutdown_system())
error = abs(numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self,
distribution,
optimizer_fn,
use_callable_loss=True):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
ds = distribution.distribute_dataset(dataset_fn)
if context.executing_eagerly():
iterator = ds.make_one_shot_iterator()
else:
iterator = ds.make_initializable_iterator()
def run_step():
return control_flow_ops.group(
distribution.unwrap(
distribution.call_for_each_replica(
model_fn,
iterator.get_next(),
run_concurrently=layer.built)))
if not context.executing_eagerly():
with self.cached_session() as sess:
sess.run(iterator.initializer)
run_step = sess.make_callable(run_step())
self.evaluate(variables.global_variables_initializer())
weights, biases = [], []
for _ in range(10):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
def testTrainNetwork(self, distribution, optimizer_fn, use_callable_loss):
with distribution.scope():
model_fn, dataset_fn, layer = minimize_loss_example(
optimizer_fn,
use_bias=True,
use_callable_loss=use_callable_loss)
def step_fn(ctx, *inputs):
del ctx # Unused
return distribution.group(
distribution.call_for_each_tower(
model_fn, *inputs, run_concurrently=layer.built))
iterator = distribution.distribute_dataset(
dataset_fn).make_one_shot_iterator()
def run_step():
return distribution.run_steps_on_dataset(step_fn,
iterator,
iterations=2).run_op
self.evaluate(distribution.initialize())
if not context.executing_eagerly():
with self.test_session() as sess:
run_step = sess.make_callable(run_step())
self.evaluate(variables_lib.global_variables_initializer())
weights, biases = [], []
for _ in range(5):
run_step()
weights.append(self.evaluate(layer.kernel))
biases.append(self.evaluate(layer.bias))
self.evaluate(distribution.finalize())
error = abs(
numpy.add(numpy.squeeze(weights), numpy.squeeze(biases)) - 1)
is_not_increasing = all(y <= x for x, y in zip(error, error[1:]))
self.assertTrue(is_not_increasing)
