def training_loop(TransformerLM, train_gen, eval_gen, output_dir="./model"):
output_dir = os.path.expanduser(output_dir)
lr_schedule = trax.lr.warmup_and_rsqrt_decay(n_warmup_steps=1000,
max_value=0.01)
# This sets up loss function and our adam optimizer used to fit the data efficiently
train_task = training.TrainTask(labeled_data=train_gen,
loss_layer=tl.CrossEntropyLoss(),
optimizer=trax.optimizers.Adam(0.01),
lr_schedule=lr_schedule,
n_steps_per_checkpoint=10)
# We evaluate on a different dataset to ensure no overfitting
eval_task = training.EvalTask(
labeled_data=eval_gen, metrics=[tl.CrossEntropyLoss(),
tl.Accuracy()])
loop = training.Loop(TransformerLM(d_model=512,
d_ff=2048,
n_layers=6,
n_heads=8,
mode='train'),
train_task,
eval_tasks=[eval_task],
output_dir=output_dir)
return loop
def test_restores_history(self):
"""Training restores history from directory where it saved it."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()])
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model, [task],
eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0,
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
loop.run(4)
loop2 = training.Loop(model, [task], output_dir=tmp_dir)
self.assertLen(loop2.history.modes, 2)
self.assertLen(loop2.history.metrics_for_mode('train'), 6)
self.assertLen(loop2.history.metrics_for_mode('eval'), 1)
for mode, metric in [
('train', 'metrics/L2Loss'),
('train', 'training/learning_rate'),
('train', 'training/steps per second'),
('train', 'training/gradients_l2'),
('train', 'training/loss'),
('train', 'training/weights_l2'),
('eval', 'metrics/L2Loss'),
]:
self.assertLen(loop2.history.get(mode, metric), 1)
self.assertEqual(2, loop2.history.get(mode, metric)[0][0])
def test_train_mnist_multitask(self, mock_stdout):
"""Train two-head MNIST model a bit, to compare to other implementations."""
mnist_model = _build_model(two_heads=True)
# MNIST classification task.
(cls_task, cls_eval_task) = _mnist_tasks(head=tl.Select([0], n_in=2))
# Auxiliary brightness prediction task.
reg_task = training.TrainTask(
itertools.cycle(_mnist_brightness_dataset().train_stream(1)),
tl.Serial(tl.Select([1]), tl.L2Loss()),
adam.Adam(0.001),
)
reg_eval_task = training.EvalTask(
itertools.cycle(_mnist_brightness_dataset().eval_stream(1)),
[tl.Serial(tl.Select([1]), tl.L2Loss())],
n_eval_batches=1,
metric_names=['L2'],
)
training_session = training.Loop(
mnist_model,
tasks=[cls_task, reg_task],
eval_tasks=[cls_eval_task, reg_eval_task],
eval_at=lambda step_n: step_n % 20 == 0,
which_task=lambda step_n: step_n % 2,
)
training_session.run(n_steps=100)
self.assertEqual(training_session.step, 100)
# Assert that we reach at least 80% eval accuracy on MNIST.
self.assertGreater(_read_metric('Accuracy', mock_stdout), 0.8)
# Assert that we get below 0.03 brightness prediction error.
self.assertLess(_read_metric('L2', mock_stdout), 0.03)
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_train_save_restore_dense(self):
"""Saves and restores a checkpoint to check for equivalence."""
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
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()
model2, training_session2 = _make_model_and_session()
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))
def test_can_predict_with_trained_model(self):
model = tl.Serial(tl.Dense(3), tl.Branch(tl.Dense(1), tl.Dense(2)))
train_tasks, eval_tasks = [], []
for output_dim in [1, 2]:
# The head we select from the model: 0 for output_dim 1 and 1 for 2.
head_index = output_dim - 1
train_tasks.append(
training.TrainTask(
_very_simple_data(output_dim),
tl.Serial(tl.Select([head_index], n_in=2), tl.L2Loss()),
optimizers.SGD(.01)))
eval_tasks.append(
training.EvalTask(
_very_simple_data(
output_dim), # deliberately re-use training data
[tl.Serial(tl.Select([head_index], n_in=2), tl.L2Loss())]))
tmp_dir = self.create_tempdir().full_path
training_session = training.Loop(
model,
tasks=train_tasks,
eval_tasks=eval_tasks,
checkpoint_at=lambda step_n: step_n == 1,
output_dir=tmp_dir,
which_task=lambda step_n: step_n % 2,
)
training_session.run(n_steps=2)
trained_model = training_session.eval_model
inp = next(_very_simple_data())[0]
out = trained_model(inp)
self.assertEqual(
shapes.signature(out),
(shapes.ShapeDtype((8, 1)), shapes.ShapeDtype((8, 2))),
)
def test_initializes_step_callbacks_with_loop_instance(self):
"""Runs a training loop, asserting that callbacks are initialized."""
class ActualLoop:
# Wrapper object to make the Loop reference mutable.
loop = None
class TestCallback(callbacks.TrainingStepCallback):
def __init__(self, loop):
super().__init__(loop)
ActualLoop.loop = loop
def call_at(self, step):
return False
def on_step_begin(self, step):
del step
def on_step_end(self, step):
del step
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
expected_loop = training.Loop(model, [task], callbacks=[TestCallback])
self.assertIs(ActualLoop.loop, expected_loop)
def test_calls_step_callbacks(self):
"""Runs a training loop, asserting that callbacks are called."""
call_at_steps = [1, 3, 4]
begin_steps = []
end_steps = []
test_case = self
class TestCallback(callbacks.TrainingStepCallback):
def call_at(self, step):
return step in call_at_steps
def on_step_begin(self, step):
begin_steps.append(step)
def on_step_end(self, step):
# Assert that on_step_begin() was called before.
test_case.assertIn(step, begin_steps)
end_steps.append(step)
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
loop = training.Loop(model, [task], callbacks=[TestCallback])
loop.run(n_steps=5)
# Assert that the callback has been called at the appropriate steps.
self.assertEqual(begin_steps, call_at_steps)
self.assertEqual(end_steps, call_at_steps)
def test_can_predict_with_trained_model(self):
model = tl.Serial(tl.Dense(3), tl.Branch(tl.Dense(1), tl.Dense(2)))
tasks = tuple(
training.TrainTask( # pylint: disable=g-complex-comprehension
_very_simple_data(output_dim),
tl.L2Loss(),
optimizers.SGD(.01),
) for output_dim in (1, 2))
eval_tasks = tuple([
training.EvalTask( # pylint: disable=g-complex-comprehension
# deliberately re-using training data
_very_simple_data(output_dim),
[tl.L2Loss()],
)
] for output_dim in (1, 2))
tmp_dir = self.create_tempdir().full_path
training_session = training.Loop(
model,
tasks=tasks,
eval_tasks=eval_tasks,
checkpoint_at=lambda step_n: step_n == 1,
output_dir=tmp_dir,
which_task=lambda step_n: step_n % 2,
)
training_session.run(n_steps=2)
trained_model = training_session.eval_model
inp = next(_very_simple_data())[0]
out = trained_model(inp)
self.assertEqual(
shapes.signature(out),
(shapes.ShapeDtype((8, 1)), shapes.ShapeDtype((8, 2))),
)
def test_train_mnist(self):
"""Train MNIST model (almost) fully, to compare to other implementations.
Evals for cross-entropy loss and accuracy are run every 50 steps;
their values are visible in the test log.
"""
mnist_model = tl.Serial(
tl.Flatten(),
tl.Dense(512),
tl.Relu(),
tl.Dense(512),
tl.Relu(),
tl.Dense(10),
tl.LogSoftmax(),
)
task = training.TrainTask(
itertools.cycle(_mnist_dataset().train_stream(1)),
tl.CrossEntropyLoss(), adafactor.Adafactor(.02))
eval_task = training.EvalTask(
itertools.cycle(_mnist_dataset().eval_stream(1)),
[tl.CrossEntropyLoss(), tl.Accuracy()],
n_eval_batches=10)
training_session = training.Loop(
mnist_model, [task],
eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 50 == 0)
training_session.run(n_steps=1000)
self.assertEqual(training_session.step, 1000)
def training_loop(n_steps=50, cutoff=0.05, output_dir="./model/"):
train_gen, eval_gen, vocab_size = generate_data(cutoff)
lr_schedule = trax.lr.warmup_and_rsqrt_decay(n_warmup_steps=1000,
max_value=0.01)
train_task = training.TrainTask(
# labeled data
labeled_data=train_gen,
# loss layer
loss_layer=tl.CrossEntropyLoss(),
# optimizer
optimizer=trax.optimizers.Adam(0.01),
# lr_schedule
lr_schedule=lr_schedule,
# n_steps
n_steps_per_checkpoint=n_steps)
eval_task = training.EvalTask(
# labeled data
labeled_data=eval_gen,
# metrics
metrics=[tl.CrossEntropyLoss(), tl.Accuracy()])
loop = training.Loop(ReformerLM(vocab_size, 6, mode='train'),
train_task,
eval_tasks=[eval_task],
output_dir=output_dir)
return loop
def test_train_mnist(self):
"""Train MNIST model (almost) fully, to compare to other implementations.
Evals for cross-entropy loss and accuracy are run every 50 steps;
their values are visible in the test log.
"""
gin.parse_config([
'batch_fn.batch_size_per_device = 256',
'batch_fn.eval_batch_size = 256',
])
mnist_model = tl.Serial(
tl.Flatten(),
tl.Dense(512),
tl.Relu(),
tl.Dense(512),
tl.Relu(),
tl.Dense(10),
tl.LogSoftmax(),
)
task = training.TrainTask(
itertools.cycle(_mnist_dataset().train_stream(1)),
tl.CrossEntropyLoss(), adafactor.Adafactor(.02))
eval_task = training.EvalTask(
itertools.cycle(_mnist_dataset().eval_stream(1)),
[tl.CrossEntropyLoss(), tl.AccuracyScalar()],
names=['CrossEntropyLoss', 'AccuracyScalar'],
eval_at=lambda step_n: step_n % 50 == 0,
eval_N=10)
training_session = training.Loop(mnist_model,
task,
eval_task=eval_task)
training_session.run(n_steps=1000)
self.assertEqual(training_session.current_step(), 1000)
def _mnist_tasks(head=None):
"""Creates MNIST training and evaluation tasks.
Args:
head: Adaptor layer to put before loss and accuracy layers in the tasks.
Returns:
A pair (train_task, eval_task) consisting of the MNIST training task and the
MNIST evaluation task using cross-entropy as loss and accuracy as metric.
"""
loss = tl.CrossEntropyLoss()
accuracy = tl.Accuracy()
if head is not None:
loss = tl.Serial(head, loss)
accuracy = tl.Serial(head, accuracy)
task = training.TrainTask(
itertools.cycle(_mnist_dataset().train_stream(1)),
loss,
adam.Adam(0.001),
)
eval_task = training.EvalTask(
itertools.cycle(_mnist_dataset().eval_stream(1)),
[loss, accuracy],
n_eval_batches=10,
metric_names=['CrossEntropy', 'Accuracy'],
)
return (task, eval_task)
def test_summaries_are_written(self):
"""Training writes down metrics when writting is turned on."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
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
training_session = training.Loop(model, [task], eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
expected_train_metric_dir = os.path.join(tmp_dir, 'train')
expected_eval_metric_dir = os.path.join(tmp_dir, 'eval')
for directory in [expected_train_metric_dir, expected_eval_metric_dir]:
self.assertFalse(
os.path.isdir(directory), 'Failed for directory %s.' % directory)
training_session.run(n_steps=15)
time.sleep(1) # wait for the files to be closed
for directory in [expected_train_metric_dir, expected_eval_metric_dir]:
self.assertTrue(
os.path.isdir(directory), 'Failed for directory %s.' % directory)
self.assertEqual(
1, _count_files(directory), 'Failed for directory %s.' % directory)
training_session.run(n_steps=5)
time.sleep(1) # wait for the files to be closed
for directory in [expected_train_metric_dir, expected_eval_metric_dir]:
self.assertEqual(
2, _count_files(directory), 'Failed for directory %s.' % directory)
def test_train_one_task_eval_two_tasks(self):
"""Trains a very simple network on one task and evaluates on two tasks."""
model = tl.Serial(tl.Dense(3), tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
export_prefix_1 = 'eval_1'
eval_task_1 = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
export_prefix=export_prefix_1,
)
export_prefix_2 = 'eval_2'
eval_task_2 = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
export_prefix=export_prefix_2,
)
training_session = training.Loop(
model,
tasks=(task, ),
eval_tasks=(eval_task_1, eval_task_2),
)
self.assertEqual(0, training_session.step)
training_session.run(n_steps=5)
self.assertEqual(5, training_session.step)
export_prefixes = [
task.export_prefix for task in training_session.eval_tasks
]
self.assertCountEqual([export_prefix_1, export_prefix_2],
export_prefixes)
def test_loop_no_eval_task(self):
"""Runs a training loop with no eval task(s)."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
training_session = training.Loop(model, [task])
# Loop should initialize and run successfully, even with no eval task.
training_session.run(n_steps=5)
def test_restores_memory_efficient_from_standard(self):
"""Training restores step from directory where it saved it."""
model = tl.Serial(tl.Dense(4), tl.Dense(1))
task_std = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.Adam(.0001))
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model, [task_std],
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
loop.run(4)
task_memeff = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.Adam)
loop2 = training.Loop(model, [task_memeff],
output_dir=tmp_dir,
use_memory_efficient_trainer=True)
loop2.run(2)
self.assertEqual(6, loop2.step)
def test_loop_no_eval_task_tfnp(self):
"""Runs a training loop with no eval task(s), TFNP backend."""
with fastmath.use_backend(fastmath.Backend.TFNP):
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.Adam(.01))
training_session = training.Loop(model, [task])
# Loop should initialize and run successfully, even with no eval task.
training_session.run(n_steps=5)
def test_restores_step(self):
"""Training restores step from directory where it saved it."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model, [task],
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
loop.run(4)
loop2 = training.Loop(model, [task], output_dir=tmp_dir)
self.assertEqual(4, loop2.step)
def test_restore_fails_different_model(self):
"""Training restores from a checkpoint created with smaller model."""
model1 = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model1, [task],
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
loop.run(2)
model2 = tl.Serial(tl.Dense(2))
with self.assertRaises(IndexError):
training.Loop(model2, [task], output_dir=tmp_dir)
def test_restores_from_smaller_model(self):
"""Training restores from a checkpoint created with smaller model."""
model1 = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.Adam(.01))
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model1, [task],
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir)
loop.run(2)
model2 = tl.Serial(tl.Dense(1), tl.Dense(1))
loop2 = training.Loop(model2, [task], output_dir=tmp_dir)
self.assertEqual(2, loop2.step)
def _mnist_tasks():
task = training.TrainTask(
itertools.cycle(_mnist_dataset().train_stream(1)),
tl.CrossEntropyLoss(),
adam.Adam(0.001),
)
eval_task = training.EvalTask(
itertools.cycle(_mnist_dataset().eval_stream(1)),
(tl.CrossEntropyLoss(), tl.Accuracy()),
n_eval_batches=10,
metric_names=('CrossEntropy', 'Accuracy'),
)
return (task, eval_task)
def test_train_dense_layer_with_momentum(self):
"""Trains with an optimizer that has slots / requires initialization."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.Momentum(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
metric_names=['Momentum.L2Loss'])
training_session = training.Loop(model, [task], eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0)
self.assertEqual(0, training_session.step)
training_session.run(n_steps=20)
self.assertEqual(20, training_session.step)
def test_restores_step_sharded_bfloat16(self):
"""Training restores step from where it saved it, sharded and bfloat16."""
model = tl.Serial(tl.Dense(1, use_bfloat16=True))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD)
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model, [task],
checkpoint_at=lambda step_n: step_n % 2 == 0,
output_dir=tmp_dir, use_memory_efficient_trainer=True)
loop.run(4)
loop2 = training.Loop(model, [task],
output_dir=tmp_dir, use_memory_efficient_trainer=True)
self.assertEqual(4, loop2.step)
loop2.run(2) # check that continued training works
self.assertEqual(6, loop2.step)
def test_train_dense_layer_evals(self):
"""Trains a very simple network on a very simple task, 2 epochs."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()])
training_session = training.Loop(model, [task], eval_tasks=[eval_task],
eval_at=lambda step_n: False)
self.assertEqual(0, training_session.step)
training_session.run(n_steps=10)
self.assertEqual(10, training_session.step)
training_session.run_evals()
self.assertEqual(10, training_session.step) # Unchanged
def test_train_dense_layer(self):
"""Trains a very simple network on a very simple task."""
model = tl.Dense(1)
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
sgd.SGD(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
names=['SGD.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_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_loop_checkpoint_high_metric(self):
"""Runs a training loop that saves checkpoints for high metric values."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(_very_simple_data(), tl.L2Loss(),
optimizers.SGD(.01))
eval_metric = tl.L2Loss()
eval_task = training.EvalTask(_very_simple_data(), [eval_metric],
metric_names=['l2_loss'])
tmp_dir = self.create_tempdir().full_path
loop = training.Loop(model, [task],
eval_tasks=[eval_task],
output_dir=tmp_dir,
eval_at=lambda step_n: step_n % 2 == 0,
checkpoint_at=lambda step_n: step_n % 2 == 0,
checkpoint_high_metric='l2_loss')
loop.run(n_steps=18)
def test_train_dense_layer(self):
"""Trains a very simple network on a very simple task."""
model = tl.Serial(tl.Dense(1))
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
metric_names=['SGD.L2Loss'])
training_session = training.Loop(model, [task], eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0)
self.assertEqual(0, training_session.step)
training_session.run(n_steps=15)
self.assertEqual(15, training_session.step)
training_session.run(n_steps=5)
self.assertEqual(20, training_session.step)
def test_loop_with_initialized_model(self):
"""Check that loop does not re-initialize an already initialized model."""
model = tl.Serial(tl.Dense(1))
example_data = next(_very_simple_data())
model.init(example_data)
w = model.weights[0][0]
task = training.TrainTask(
_very_simple_data(), tl.L2Loss(), optimizers.SGD(.01))
eval_task = training.EvalTask(
_very_simple_data(), # deliberately re-using training data
[tl.L2Loss()],
metric_names=['SGD.L2Loss'])
loop = training.Loop(model, [task], eval_tasks=[eval_task],
eval_at=lambda step_n: step_n % 2 == 0)
self.assertEqual(0, loop.step)
self.assertEqual(loop.model.weights[0][0], w)
