trainer inheritance and test

Author: awaelchli

pytorch_lightning/trainer/callback_hook.py

@@ -6,11 +6,10 @@
 
 class TrainerCallbackHookMixin(ABC):
 
-    def __init__(self):
-        # this is just a summary on variables used in this abstract class,
-        # the proper values/initialisation should be done in child class
-        self.callbacks: List[Callback] = []
-        self.get_model: Callable = ...
+    # this is just a summary on variables used in this abstract class,
+    # the proper values/initialisation should be done in child class
+    callbacks: List[Callback] = []
+    get_model: Callable = ...
 
     def on_init_start(self):
         """Called when the trainer initialization begins, model has not yet been set."""

pytorch_lightning/trainer/trainer.py

@@ -304,6 +304,7 @@ def __init__(
                 Additionally, can be set to either `power` that estimates the batch size through
                 a power search or `binsearch` that estimates the batch size through a binary search.
         """
+        super().__init__()
 
         self.deterministic = deterministic
         torch.backends.cudnn.deterministic = self.deterministic

tests/trainer/test_trainer.py

@@ -772,3 +772,39 @@ def test_trainer_config(trainer_kwargs, expected):
     assert trainer.on_gpu is expected["on_gpu"]
     assert trainer.single_gpu is expected["single_gpu"]
     assert trainer.num_processes == expected["num_processes"]
+
+
+def test_trainer_subclassing():
+
+    model = EvalModelTemplate()
+
+    # First way of pulling out args from signature is to list them
+    class TrainerSubclass(Trainer):
+
+        def __init__(self, custom_arg, *args, custom_kwarg='test', **kwargs):
+            super().__init__(*args, **kwargs)
+            self.custom_arg = custom_arg
+            self.custom_kwarg = custom_kwarg
+
+    trainer = TrainerSubclass(123, custom_kwarg='custom', fast_dev_run=True)
+    result = trainer.fit(model)
+    assert result == 1
+    assert trainer.custom_arg == 123
+    assert trainer.custom_kwarg == 'custom'
+    assert trainer.fast_dev_run
+
+    # Second way is to pop from the dict
+    # It's a special case because Trainer does not have any positional args
+    class TrainerSubclass(Trainer):
+
+        def __init__(self, **kwargs):
+            self.custom_arg = kwargs.pop('custom_arg', 0)
+            self.custom_kwarg = kwargs.pop('custom_kwarg', 'test')
+            super().__init__(**kwargs)
+
+    trainer = TrainerSubclass(custom_kwarg='custom', fast_dev_run=True)
+    result = trainer.fit(model)
+    assert result == 1
+    assert trainer.custom_kwarg == 'custom'
+    assert trainer.fast_dev_run
+