Update evaluation hook tests

tests/callbacks/test_callbacks.py

@@ -100,76 +100,6 @@ def test_trainer_callback_hook_system_fit(_, tmpdir):
     ]
 
 
-def test_trainer_callback_hook_system_test(tmpdir):
-    """Test the callback hook system for test."""
-
-    model = BoringModel()
-    callback_mock = MagicMock()
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        callbacks=[callback_mock],
-        max_epochs=1,
-        limit_test_batches=2,
-        progress_bar_refresh_rate=0,
-    )
-
-    trainer.test(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, stage='test'),
-        call.on_configure_sharded_model(trainer, model),
-        call.on_test_start(trainer, model),
-        call.on_epoch_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_epoch_end(trainer, model),
-        call.on_test_end(trainer, model),
-        call.teardown(trainer, model, stage='test'),
-    ]
-
-
-def test_trainer_callback_hook_system_validate(tmpdir):
-    """Test the callback hook system for validate."""
-
-    model = BoringModel()
-    callback_mock = MagicMock()
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        callbacks=[callback_mock],
-        max_epochs=1,
-        limit_val_batches=2,
-        progress_bar_refresh_rate=0,
-    )
-
-    trainer.validate(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, stage='validate'),
-        call.on_configure_sharded_model(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_batch_start(trainer, model, ANY, 1, 0),
-        call.on_validation_batch_end(trainer, model, ANY, ANY, 1, 0),
-        call.on_validation_epoch_end(trainer, model),
-        call.on_epoch_end(trainer, model),
-        call.on_validation_end(trainer, model),
-        call.teardown(trainer, model, stage='validate'),
-    ]
-
-
 def test_callbacks_configured_in_model(tmpdir):
     """ Test the callback system with callbacks added through the model hook. """
 

tests/models/test_hooks.py

@@ -314,18 +314,41 @@ def _val_batch():
         ]
 
     @staticmethod
-    def _test_batch():
+    def _eval_epoch(fn, trainer, model, batches, key):
+        outputs = {key: ANY}
         return [
-            'on_test_batch_start',
-            'on_before_batch_transfer',
-            'transfer_batch_to_device',
-            'on_after_batch_transfer',
-            'forward',
-            'test_step',
-            'test_step_end',
-            'on_test_batch_end',
+            dict(name='Callback.on_epoch_start', args=(trainer, model)),
+            dict(name='on_epoch_start'),
+            dict(name=f'Callback.on_{fn}_epoch_start', args=(trainer, model)),
+            dict(name=f'on_{fn}_epoch_start'),
+            *HookedModel._eval_batch(fn, trainer, model, batches, key),
+            dict(name=f'{fn}_epoch_end', args=([outputs] * batches, )),
+            dict(name=f'Callback.on_{fn}_epoch_end', args=(trainer, model)),
+            dict(name=f'on_{fn}_epoch_end'),
+            dict(name='Callback.on_epoch_end', args=(trainer, model)),
+            dict(name='on_epoch_end'),
         ]
 
+    @staticmethod
+    def _eval_batch(fn, trainer, model, batches, key):
+        out = []
+        outputs = {key: ANY}
+        for i in range(batches):
+            out.extend([
+                # TODO: `{,Callback}.on_batch_{start,end}`
+                dict(name=f'Callback.on_{fn}_batch_start', args=(trainer, model, ANY, i, 0)),
+                dict(name=f'on_{fn}_batch_start', args=(ANY, i, 0)),
+                dict(name='on_before_batch_transfer', args=(ANY, None)),
+                dict(name='transfer_batch_to_device', args=(ANY, torch.device('cpu'), None)),
+                dict(name='on_after_batch_transfer', args=(ANY, None)),
+                dict(name='forward', args=(ANY, )),
+                dict(name=f'{fn}_step', args=(ANY, i)),
+                dict(name=f'{fn}_step_end', args=(outputs, )),
+                dict(name=f'Callback.on_{fn}_batch_end', args=(trainer, model, outputs, ANY, i, 0)),
+                dict(name=f'on_{fn}_batch_end', args=(outputs, ANY, i, 0)),
+            ])
+        return out
+
     @staticmethod
     def _predict_batch(trainer, model, batches):
         out = []
@@ -464,83 +487,57 @@ def test_trainer_model_hook_system_fit_no_val(tmpdir):
 
 
 @pytest.mark.parametrize('batches', (0, 2))
-def test_trainer_model_hook_system_validate(tmpdir, batches):
+@pytest.mark.parametrize(['verb', 'noun', 'dataloader', 'key'], [
+    ('validate', 'validation', 'val', 'x'),
+    ('test', 'test', 'test', 'y'),
+])
+def test_trainer_model_hook_system_eval(tmpdir, batches, verb, noun, dataloader, key):
     called = []
     model = HookedModel(called)
+    callback = HookedCallback(called)
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=1,
         limit_val_batches=batches,
+        limit_test_batches=batches,
         progress_bar_refresh_rate=0,
         weights_summary=None,
+        callbacks=[callback],
     )
-    assert called == []
-    trainer.validate(model, verbose=False)
+    assert called == [
+        dict(name='Callback.on_init_start', args=(trainer, )),
+        dict(name='Callback.on_init_end', args=(trainer, )),
+    ]
+    fn = getattr(trainer, verb)
+    fn(model, verbose=False)
     hooks = [
-        'train',  # eval() == train(False)
-        'on_validation_model_eval',
-        'zero_grad',
-        'on_validation_start',
-        'on_epoch_start',
-        'on_validation_epoch_start',
-        *(HookedModel._val_batch() * batches),
-        'validation_epoch_end',
-        'on_validation_epoch_end',
-        'on_epoch_end',
-        'on_validation_end',
-        'train',
-        'on_validation_model_train'
+        dict(name='train', args=(False, )),
+        dict(name=f'on_{noun}_model_eval'),
+        dict(name='zero_grad'),
+        dict(name=f'Callback.on_{noun}_start', args=(trainer, model)),
+        dict(name=f'on_{noun}_start'),
+        *model._eval_epoch(noun, trainer, model, batches, key),
+        dict(name=f'Callback.on_{noun}_end', args=(trainer, model)),
+        dict(name=f'on_{noun}_end'),
+        dict(name='train'),
+        dict(name=f'on_{noun}_model_train'),
     ]
     expected = [
-        'prepare_data',
-        'configure_callbacks',
-        'setup',
-        'configure_sharded_model',
-        'on_val_dataloader',
-        'val_dataloader',
+        dict(name='Callback.on_init_start', args=(trainer, )),
+        dict(name='Callback.on_init_end', args=(trainer, )),
+        dict(name='prepare_data'),
+        dict(name='configure_callbacks'),
+        dict(name='Callback.on_before_accelerator_backend_setup', args=(trainer, model)),
+        dict(name='Callback.setup', args=(trainer, model), kwargs=dict(stage=verb)),
+        dict(name='setup', kwargs=dict(stage=verb)),
+        dict(name='configure_sharded_model'),
+        dict(name='Callback.on_configure_sharded_model', args=(trainer, model)),
+        dict(name=f'on_{dataloader}_dataloader'),
+        dict(name=f'{dataloader}_dataloader'),
         *(hooks if batches else []),
-        'teardown',
+        dict(name='Callback.teardown', args=(trainer, model), kwargs=dict(stage=verb)),
+        dict(name='teardown', kwargs=dict(stage=verb)),
     ]
-    called = [c['name'] for c in called]
-    assert called == expected
-
-
-def test_trainer_model_hook_system_test(tmpdir):
-    called = []
-    model = HookedModel(called)
-    batches = 2
-    trainer = Trainer(
-        default_root_dir=tmpdir,
-        max_epochs=1,
-        limit_test_batches=batches,
-        progress_bar_refresh_rate=0,
-        weights_summary=None,
-    )
-    assert called == []
-    trainer.test(model, verbose=False)
-    expected = [
-        'prepare_data',
-        'configure_callbacks',
-        'setup',
-        'configure_sharded_model',
-        'on_test_dataloader',
-        'test_dataloader',
-        'train',  # eval() == train(False)
-        'on_test_model_eval',
-        'zero_grad',
-        'on_test_start',
-        'on_epoch_start',
-        'on_test_epoch_start',
-        *(HookedModel._test_batch() * batches),
-        'test_epoch_end',
-        'on_test_epoch_end',
-        'on_epoch_end',
-        'on_test_end',
-        'train',
-        'on_test_model_train',
-        'teardown',
-    ]
-    called = [c['name'] for c in called]
     assert called == expected