def test_executes_with(self, spec, momentum):
weights = tf.nest.map_structure(lambda s: tf.ones(s.shape, s.dtype), spec)
gradients = tf.nest.map_structure(lambda s: tf.ones(s.shape, s.dtype), spec)
optimizer = sgdm.SGD(0.01, momentum=momentum)
state = optimizer.initialize(spec)
for _ in range(10):
state, weights = optimizer.next(state, weights, gradients)
tf.nest.map_structure(lambda w: self.assertTrue(all(tf.math.is_finite(w))),
weights)
def test_convergence(self, momentum):
weights, fn, grad_fn = optimizer_test_utils.test_quadratic_problem()
self.assertGreater(fn(weights), 5.0)
optimizer = sgdm.SGD(0.1, momentum=momentum)
state = optimizer.initialize(tf.TensorSpec(weights.shape, weights.dtype))
for _ in range(100):
gradients = grad_fn(weights)
state, weights = optimizer.next(state, weights, gradients)
self.assertLess(fn(weights), 0.005)
def test_math_momentum_0_5(self):
weights = tf.constant([1.0], tf.float32)
gradients = tf.constant([2.0], tf.float32)
optimizer = sgdm.SGD(0.01, momentum=0.5)
history = [weights]
state = optimizer.initialize(_SCALAR_SPEC)
for _ in range(4):
state, weights = optimizer.next(state, weights, gradients)
history.append(weights)
self.assertAllClose([[1.0], [0.98], [0.95], [0.915], [0.8775]], history)
def test_math_no_momentum(self):
weights = tf.constant([1.0], tf.float32)
gradients = tf.constant([2.0], tf.float32)
optimizer = sgdm.SGD(0.01)
history = [weights]
state = optimizer.initialize(_SCALAR_SPEC)
self.assertTupleEqual((), state)
for _ in range(4):
state, weights = optimizer.next(state, weights, gradients)
history.append(weights)
self.assertAllClose([[1.0], [0.98], [0.96], [0.94], [0.92]], history)