def test_train_dense_layer_with_momentum(self):
"""Trains with an optimizer that has slots / requires initialization."""
model = tl.Dense(1)
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
momentum.Momentum(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
names=['Momentum.L2Loss'],
eval_at=lambda step_n: step_n % 2 == 0,
eval_N=1)
training_session = training.Loop(model, task, eval_task=eval_task)
self.assertIsNone(training_session.current_step())
training_session.run(n_steps=20)
self.assertEqual(20, training_session.current_step())
def test_slots(self):
weights_shape = (3, 5)
weight_tree = np.arange(15).reshape(weights_shape)
# SGD - an optimizer that doesn't use slots.
opt_1 = optimizers.SGD(.01)
self.assertIsNone(opt_1.slots)
opt_1.tree_init(weight_tree)
self.assertIsInstance(opt_1.slots, tuple)
self.assertLen(opt_1.slots, 1)
self.assertIsNone(opt_1.slots[0])
# Momentum - an optimizer with slots
opt_2 = momentum.Momentum(.01)
self.assertIsNone(opt_2.slots)
opt_2.tree_init(weight_tree)
self.assertIsInstance(opt_2.slots, tuple)
self.assertLen(opt_2.slots, 1)
self.assertEqual(weights_shape, opt_2.slots[0].shape)