def _test_fitting(self, params, logs, has_valid=True):
self.assertEqual(len(logs), params['epochs'])
train_dict = dict(zip(self.metrics_names, self.metrics_values), loss=ANY)
if has_valid:
val_metrics_names = ['val_' + metric_name for metric_name in self.metrics_names]
val_dict = dict(zip(val_metrics_names, self.metrics_values), val_loss=ANY)
log_dict = {**train_dict, **val_dict}
else:
log_dict = train_dict
for epoch, log in enumerate(logs, 1):
self.assertEqual(log, dict(log_dict, epoch=epoch))
call_list = []
call_list.append(call.on_train_begin({}))
for epoch in range(1, params['epochs']+1):
call_list.append(call.on_epoch_begin(epoch, {}))
for step in range(1, params['steps']+1):
call_list.append(call.on_batch_begin(step, {}))
call_list.append(call.on_backward_end(step))
call_list.append(call.on_batch_end(step, {'batch': step, 'size': ANY, **train_dict}))
call_list.append(call.on_epoch_end(epoch, {'epoch': epoch, **log_dict}))
call_list.append(call.on_train_end({}))
method_calls = self.mock_callback.method_calls
self.assertIn(call.set_model(self.model), method_calls[:2])
self.assertIn(call.set_params(params), method_calls[:2])
self.assertEqual(len(method_calls), len(call_list) + 2)
self.assertEqual(method_calls[2:], call_list)
def _test_batch_delay(self, epoch_delay, batch_in_epoch_delay):
batch_delay = epoch_delay*DelayCallbackTest.steps_per_epoch + batch_in_epoch_delay
delay_callback = DelayCallback(self.mock_callback, batch_delay=batch_delay)
train_generator = some_data_generator(DelayCallbackTest.batch_size)
valid_generator = some_data_generator(DelayCallbackTest.batch_size)
logs = self.model.fit_generator(train_generator, valid_generator, epochs=DelayCallbackTest.epochs, steps_per_epoch=DelayCallbackTest.steps_per_epoch, validation_steps=DelayCallbackTest.steps_per_epoch, callbacks=[delay_callback])
params = {'epochs': DelayCallbackTest.epochs, 'steps': DelayCallbackTest.steps_per_epoch}
call_list = []
call_list.append(call.on_train_begin({}))
for epoch in range(epoch_delay + 1, DelayCallbackTest.epochs+1):
call_list.append(call.on_epoch_begin(epoch, {}))
start_step = batch_in_epoch_delay + 1 if epoch == epoch_delay + 1 else 1
for step in range(start_step, params['steps']+1):
call_list.append(call.on_batch_begin(step, {}))
call_list.append(call.on_backward_end(step))
call_list.append(call.on_batch_end(step, {'batch': step, 'size': DelayCallbackTest.batch_size, **self.train_dict}))
call_list.append(call.on_epoch_end(epoch, {'epoch': epoch, **self.log_dict}))
call_list.append(call.on_train_end({}))
method_calls = self.mock_callback.method_calls
self.assertIn(call.set_model(self.model), method_calls[:2])
self.assertIn(call.set_params(params), method_calls[:2])
self.assertEqual(len(method_calls), len(call_list) + 2)
self.assertEqual(method_calls[2:], call_list)
def test_trainer_callback_system(torch_save):
"""Test the callback system."""
model = BoringModel()
callback_mock = MagicMock()
trainer_options = dict(
callbacks=[callback_mock],
max_epochs=1,
limit_val_batches=1,
limit_train_batches=3,
limit_test_batches=2,
progress_bar_refresh_rate=0,
)
# no call yet
callback_mock.assert_not_called()
# fit model
trainer = Trainer(**trainer_options)
# check that only the to calls exists
assert trainer.callbacks[0] == callback_mock
assert callback_mock.method_calls == [
call.on_init_start(trainer),
call.on_init_end(trainer),
]
trainer.fit(model)
assert callback_mock.method_calls == [
call.on_init_start(trainer),
call.on_init_end(trainer),
call.setup(trainer, model, 'fit'),
call.on_fit_start(trainer, model),
call.on_pretrain_routine_start(trainer, model),
call.on_pretrain_routine_end(trainer, model),
call.on_sanity_check_start(trainer, model),
call.on_validation_start(trainer, model),
call.on_validation_epoch_start(trainer, model),
call.on_validation_batch_start(trainer, model, ANY, 0, 0),
call.on_validation_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_validation_epoch_end(trainer, model),
call.on_validation_end(trainer, model),
call.on_sanity_check_end(trainer, model),
call.on_train_start(trainer, model),
call.on_epoch_start(trainer, model),
call.on_train_epoch_start(trainer, model),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 0, 0),
call.on_after_backward(trainer, model),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_batch_end(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 1, 0),
call.on_after_backward(trainer, model),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_batch_end(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 1, 0),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 2, 0),
call.on_after_backward(trainer, model),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_batch_end(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 2, 0),
call.on_validation_start(trainer, model),
call.on_validation_epoch_start(trainer, model),
call.on_validation_batch_start(trainer, model, ANY, 0, 0),
call.on_validation_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_validation_epoch_end(trainer, model),
call.on_validation_end(trainer, model),
call.on_save_checkpoint(trainer, model),
call.on_epoch_end(trainer, model),
call.on_train_epoch_end(trainer, model, ANY),
call.on_train_end(trainer, model),
call.on_fit_end(trainer, model),
call.teardown(trainer, model, 'fit'),
]
callback_mock.reset_mock()
trainer = Trainer(**trainer_options)
trainer.test(model)
assert callback_mock.method_calls == [
call.on_init_start(trainer),
call.on_init_end(trainer),
call.setup(trainer, model, 'test'),
call.on_fit_start(trainer, model),
call.on_test_start(trainer, model),
call.on_test_epoch_start(trainer, model),
call.on_test_batch_start(trainer, model, ANY, 0, 0),
call.on_test_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_test_batch_start(trainer, model, ANY, 1, 0),
call.on_test_batch_end(trainer, model, ANY, ANY, 1, 0),
call.on_test_epoch_end(trainer, model),
call.on_test_end(trainer, model),
call.on_fit_end(trainer, model),
call.teardown(trainer, model, 'fit'),
call.teardown(trainer, model, 'test'),
]
def test_trainer_callback_hook_system_fit(_, tmpdir):
"""Test the callback hook system for fit."""
model = BoringModel()
callback_mock = MagicMock()
trainer = Trainer(
default_root_dir=tmpdir,
callbacks=[callback_mock],
max_epochs=1,
limit_val_batches=1,
limit_train_batches=3,
progress_bar_refresh_rate=0,
)
# check that only the to calls exists
assert trainer.callbacks[0] == callback_mock
assert callback_mock.method_calls == [
call.on_init_start(trainer),
call.on_init_end(trainer),
]
# fit model
trainer.fit(model)
assert callback_mock.method_calls == [
call.on_init_start(trainer),
call.on_init_end(trainer),
call.on_before_accelerator_backend_setup(trainer, model),
call.setup(trainer, model, 'fit'),
call.on_configure_sharded_model(trainer, model),
call.on_fit_start(trainer, model),
call.on_pretrain_routine_start(trainer, model),
call.on_pretrain_routine_end(trainer, model),
call.on_sanity_check_start(trainer, model),
call.on_validation_start(trainer, model),
call.on_epoch_start(trainer, model),
call.on_validation_epoch_start(trainer, model),
call.on_validation_batch_start(trainer, model, ANY, 0, 0),
call.on_validation_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_validation_epoch_end(trainer, model),
call.on_epoch_end(trainer, model),
call.on_validation_end(trainer, model),
call.on_sanity_check_end(trainer, model),
call.on_train_start(trainer, model),
call.on_epoch_start(trainer, model),
call.on_train_epoch_start(trainer, model),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 0, 0),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_after_backward(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_batch_end(trainer, model),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 1, 0),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_after_backward(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 1, 0),
call.on_batch_end(trainer, model),
call.on_batch_start(trainer, model),
call.on_train_batch_start(trainer, model, ANY, 2, 0),
call.on_before_zero_grad(trainer, model, trainer.optimizers[0]),
call.on_after_backward(trainer, model),
call.on_train_batch_end(trainer, model, ANY, ANY, 2, 0),
call.on_batch_end(trainer, model),
call.on_train_epoch_end(trainer, model, ANY),
call.on_epoch_end(trainer, model),
call.on_validation_start(trainer, model),
call.on_epoch_start(trainer, model),
call.on_validation_epoch_start(trainer, model),
call.on_validation_batch_start(trainer, model, ANY, 0, 0),
call.on_validation_batch_end(trainer, model, ANY, ANY, 0, 0),
call.on_validation_epoch_end(trainer, model),
call.on_epoch_end(trainer, model),
call.on_validation_end(trainer, model),
call.on_save_checkpoint(
trainer,
model), # should take ANY but we are inspecting signature for BC
call.on_train_end(trainer, model),
call.on_fit_end(trainer, model),
call.teardown(trainer, model, 'fit'),
]
