def test_train_save_restore_sharded(self):
"""Saves and restores a sharded checkpoint to check for equivalence."""
if fastmath.local_device_count() < 2:
return # multi-accelerator only
base.N_WEIGHTS_SHARDS = fastmath.local_device_count()
train_data = data.Serial(lambda _: _very_simple_data(2, 2),
data.CountAndSkip('simple_data'))
task = training.TrainTask(train_data(), tl.L2Loss(),
optimizers.Adam(.0001))
eval_task = training.EvalTask(
_very_simple_data(2, 2), # deliberately re-using training data
[tl.L2Loss()],
metric_names=['SGD.L2Loss'])
tmp_dir = self.create_tempdir().full_path
def _make_model_and_session():
m = tl.Serial(tl.Dense(2))
ts = training.Loop(m, [task],
eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
return m, ts
_, training_session = _make_model_and_session()
self.assertEqual(0, training_session.step)
training_session.run(n_steps=1)
training_session.save_checkpoint('model')
_, training_session2 = _make_model_and_session()
training_session2.run(n_steps=1)
base.N_WEIGHTS_SHARDS = 1
def test_count_and_skip(self):
dataset = lambda _: ((i, i + 1) for i in range(10))
examples = data.Serial(dataset, data.CountAndSkip('toy_data'))
ex_generator = examples()
ex1 = next(ex_generator)
self.assertEqual(ex1, (0, 1))
self.assertEqual(data.inputs.data_counters['toy_data'], 1)
ex2 = next(ex_generator)
self.assertEqual(ex2, (1, 2))
self.assertEqual(data.inputs.data_counters['toy_data'], 2)
ex3 = next(examples()) # new generator, will skip
self.assertEqual(ex3, (2, 3))
self.assertEqual(data.inputs.data_counters['toy_data'], 3)
data.inputs.data_counters['toy_data'] = 0 # reset
ex4 = next(examples()) # new generator, was reset
self.assertEqual(ex4, (0, 1))
self.assertEqual(data.inputs.data_counters['toy_data'], 1)
def test_train_save_restore_dense(self):
"""Saves and restores a checkpoint to check for equivalence."""
train_data = data.Serial(lambda _: _very_simple_data(),
data.CountAndSkip('simple_data'))
task = training.TrainTask(train_data(), tl.L2Loss(),
optimizers.Adam(.0001))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
metric_names=['SGD.L2Loss'])
tmp_dir = self.create_tempdir().full_path
def _make_model_and_session():
m = tl.Serial(tl.Dense(1))
ts = training.Loop(m, [task],
eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
return m, ts
model, training_session = _make_model_and_session()
self.assertEqual(0, training_session.step)
training_session.run(n_steps=1)
training_session.save_checkpoint()
self.assertEqual(data.inputs.data_counters['simple_data'], 2)
data.inputs.data_counters['simple_data'] = 0 # reset manually
self.assertEqual(data.inputs.data_counters['simple_data'], 0) # check
model2, training_session2 = _make_model_and_session()
self.assertEqual(data.inputs.data_counters['simple_data'],
2) # restored
x = np.ones((8, 1))
y1 = model(x, rng=fastmath.random.get_prng(0))
y2 = model2(x, rng=fastmath.random.get_prng(0))
self.assertEqual(str(y1), str(y2))
training_session2.run(n_steps=1)
y1 = model(x, rng=fastmath.random.get_prng(0))
y2 = model2(x, rng=fastmath.random.get_prng(0))
self.assertNotEqual(str(y1), str(y2))
slots1 = training_session._trainer_per_task[0].slots
slots2 = training_session2._trainer_per_task[0].slots
np.testing.assert_array_equal(slots1, slots2)