def test_spec(self, specs, disjoint_init_and_next):
"""Test compatibility with different structures of variables."""
optimizer_fn = lambda: tf.keras.optimizers.SGD(0.1)
variables = tf.nest.map_structure(
lambda s: tf.Variable(tf.ones(s.shape, s.dtype)), specs)
gradients = tf.nest.map_structure(lambda s: tf.ones(s.shape, s.dtype),
specs)
optimizer = keras_optimizer.KerasOptimizer(
optimizer_fn, variables, disjoint_init_and_next=disjoint_init_and_next)
state = optimizer.initialize(specs)
for _ in range(3):
state, variables = optimizer.next(state, variables, gradients)
expected_variables = tf.nest.map_structure(
lambda s: 0.7 * tf.ones(s.shape, s.dtype), specs)
self.assertAllClose(expected_variables, variables)
def next_fn(state, initial_weights):
variables = tf.Variable(initial_weights)
optimizer = keras_optimizer.KerasOptimizer(
optimizer_fn, variables, disjoint_init_and_next=True)
return single_step(optimizer, state, variables)
def initialize_fn():
variables = tf.Variable(tf.zeros([5, 1]))
optimizer = keras_optimizer.KerasOptimizer(
optimizer_fn, variables, disjoint_init_and_next=True)
return optimizer.initialize(
tf.TensorSpec(variables.shape, variables.dtype))
def test_computation(initial_weights):
variables = tf.Variable(initial_weights)
optimizer = keras_optimizer.KerasOptimizer(
optimizer_fn, variables, disjoint_init_and_next=False)
return training_loop(optimizer, variables)