ref: modular is_overridden (Lightning-AI#3290)

* ref: modular is_overridden

* ref: modular is_overridden

* ref: modular is_overridden

* ref: modular is_overridden

Author: williamFalcon

pytorch_lightning/trainer/callback_config.py

@@ -18,6 +18,8 @@
 from pytorch_lightning.callbacks import Callback, ModelCheckpoint, EarlyStopping, ProgressBarBase, ProgressBar
 from pytorch_lightning.loggers import LightningLoggerBase
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.model_utils import is_overridden
+from pytorch_lightning.core.lightning import LightningModule
 
 
 class TrainerCallbackConfigMixin(ABC):
@@ -41,13 +43,13 @@ def save_checkpoint(self, *args):
         """Warning: this is just empty shell for code implemented in other class."""
 
     @abstractmethod
-    def is_overridden(self, *args):
+    def get_model(self) -> LightningModule:
         """Warning: this is just empty shell for code implemented in other class."""
 
     def configure_checkpoint_callback(self, checkpoint_callback):
         if checkpoint_callback is True:
             # when no val step is defined, use 'loss' otherwise 'val_loss'
-            train_step_only = not self.is_overridden('validation_step')
+            train_step_only = not is_overridden('validation_step', self.get_model())
             monitor_key = 'loss' if train_step_only else 'val_loss'
             checkpoint_callback = ModelCheckpoint(
                 filepath=None,

pytorch_lightning/trainer/data_loading.py

@@ -26,6 +26,7 @@
 from pytorch_lightning.utilities.data import has_iterable_dataset, has_len
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.debugging import InternalDebugger
+from pytorch_lightning.utilities.model_utils import is_overridden
 
 
 try:
@@ -78,10 +79,6 @@ class TrainerDataLoadingMixin(ABC):
     distributed_backend: Optional[str]
     dev_debugger: InternalDebugger
 
-    @abstractmethod
-    def is_overridden(self, *args):
-        """Warning: this is just empty shell for code implemented in other class."""
-
     def _worker_check(self, dataloader: DataLoader, name: str) -> None:
         on_windows = platform.system() == 'Windows'
 
@@ -305,8 +302,8 @@ def reset_val_dataloader(self, model: LightningModule) -> None:
         Args:
             model: The current `LightningModule`
         """
-        has_loader = self.is_overridden('val_dataloader', model)
-        has_step = self.is_overridden('validation_step', model)
+        has_loader = is_overridden('val_dataloader', model)
+        has_step = is_overridden('validation_step', model)
         if has_loader and has_step:
             self.num_val_batches, self.val_dataloaders = self._reset_eval_dataloader(model, 'val')
 
@@ -316,8 +313,8 @@ def reset_test_dataloader(self, model) -> None:
         Args:
             model: The current `LightningModule`
         """
-        has_loader = self.is_overridden('test_dataloader', model)
-        has_step = self.is_overridden('test_step', model)
+        has_loader = is_overridden('test_dataloader', model)
+        has_step = is_overridden('test_step', model)
         if has_loader and has_step:
             self.num_test_batches, self.test_dataloaders =\
                 self._reset_eval_dataloader(model, 'test')

pytorch_lightning/trainer/evaluate_loop.py

@@ -3,6 +3,7 @@
 from pytorch_lightning.core.step_result import Result, EvalResult
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities import flatten_dict
+from pytorch_lightning.utilities.model_utils import is_overridden
 
 
 class EvaluationLoop(object):
@@ -179,15 +180,15 @@ def __run_eval_epoch_end(self, num_dataloaders, using_eval_result):
         user_reduced = False
 
         if self.testing:
-            if self.trainer.is_overridden('test_epoch_end', model=model):
+            if is_overridden('test_epoch_end', model=model):
                 if using_eval_result:
                     eval_results = self.__gather_epoch_end_eval_results(outputs)
 
                 eval_results = model.test_epoch_end(eval_results)
                 user_reduced = True
 
         else:
-            if self.trainer.is_overridden('validation_epoch_end', model=model):
+            if is_overridden('validation_epoch_end', model=model):
                 if using_eval_result:
                     eval_results = self.__gather_epoch_end_eval_results(outputs)
 

pytorch_lightning/trainer/evaluation_loop.py

@@ -201,10 +201,6 @@ def copy_trainer_model_properties(self, *args):
     def get_model(self) -> LightningModule:
         """Warning: this is just empty shell for code implemented in other class."""
 
-    @abstractmethod
-    def is_overridden(self, *args):
-        """Warning: this is just empty shell for code implemented in other class."""
-
     @abstractmethod
     def transfer_batch_to_gpu(self, *args):
         """Warning: this is just empty shell for code implemented in other class."""

pytorch_lightning/trainer/model_hooks.py

@@ -26,34 +26,6 @@ def is_function_implemented(self, f_name, model=None):
         f_op = getattr(model, f_name, None)
         return callable(f_op)
 
-    def is_overridden(self, method_name: str, model: LightningModule = None) -> bool:
-        if model is None:
-            model = self.get_model()
-        # if you pass DataModule instead of None or a LightningModule, we use LightningDataModule as super
-        # TODO - refector this function to accept model_name, instance, parent so it makes more sense
-        super_object = LightningModule if not isinstance(model, LightningDataModule) else LightningDataModule
-
-        # assert model, 'no model passes'
-
-        if not hasattr(model, method_name):
-            # in case of calling deprecated method
-            return False
-
-        instance_attr = getattr(model, method_name)
-        if not instance_attr:
-            return False
-        super_attr = getattr(super_object, method_name)
-
-        # when code pointers are different, it was implemented
-        if hasattr(instance_attr, 'patch_loader_code'):
-            # cannot pickle __code__ so cannot verify if PatchDataloader
-            # exists which shows dataloader methods have been overwritten.
-            # so, we hack it by using the string representation
-            is_overridden = instance_attr.patch_loader_code != str(super_attr.__code__)
-        else:
-            is_overridden = instance_attr.__code__ is not super_attr.__code__
-        return is_overridden
-
     def has_arg(self, f_name, arg_name):
         model = self.get_model()
         f_op = getattr(model, f_name, None)

pytorch_lightning/trainer/trainer.py

@@ -56,6 +56,7 @@
 from pytorch_lightning.utilities.cloud_io import is_remote_path
 from pytorch_lightning.trainer.evaluate_loop import EvaluationLoop
 from pytorch_lightning.trainer.data_connector import DataConnector
+from pytorch_lightning.utilities.model_utils import is_overridden
 
 # warnings to ignore in trainer
 warnings.filterwarnings(
@@ -902,7 +903,7 @@ def disable_validation(self) -> bool:
     @property
     def enable_validation(self) -> bool:
         """ Check if we should run validation during training. """
-        val_loop_enabled = self.is_overridden('validation_step') and self.limit_val_batches > 0
+        val_loop_enabled = is_overridden('validation_step', self.get_model()) and self.limit_val_batches > 0
         return val_loop_enabled or self.fast_dev_run
 
     @property
@@ -1091,7 +1092,7 @@ def select_accelerator(self):
 
     def can_prepare_data(self):
         should_call_dm_prepare_data = True
-        if self.datamodule is not None and self.is_overridden('prepare_data', self.datamodule):
+        if self.datamodule is not None and is_overridden('prepare_data', self.datamodule):
             should_call_dm_prepare_data = not self.datamodule.has_prepared_data
 
         if self.prepare_data_per_node:
@@ -1197,7 +1198,7 @@ def run_pretrain_routine(self, model: LightningModule):
         self.train()
 
     def _run_sanity_check(self, ref_model, model):
-        using_val_step = ref_model.val_dataloader is not None and self.is_overridden('validation_step')
+        using_val_step = ref_model.val_dataloader is not None and is_overridden('validation_step', self.get_model())
         should_sanity_check = using_val_step and self.num_sanity_val_steps > 0 and self.limit_val_batches > 0
 
         # run tiny validation (if validation defined)
@@ -1418,8 +1419,8 @@ def call_hook(self, hook_name, *args, **kwargs):
 
             # next call hook in lightningModule
             output = None
-            if self.is_overridden(hook_name):
-                model_ref = self.get_model()
+            model_ref = self.get_model()
+            if is_overridden(hook_name, model_ref):
                 hook_fx = getattr(model_ref, hook_name)
                 output = hook_fx(*args, **kwargs)
 

pytorch_lightning/trainer/training_loop.py

@@ -181,6 +181,7 @@ def training_step(self, batch, batch_idx):
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.memory import recursive_detach
 from pytorch_lightning.utilities.parsing import AttributeDict
+from pytorch_lightning.utilities.model_utils import is_overridden
 
 try:
     from apex import amp
@@ -300,10 +301,6 @@ def clip_gradients(self, *args):
     def detect_nan_tensors(self, *args):
         """Warning: this is just empty shell for code implemented in other class."""
 
-    @abstractmethod
-    def is_overridden(self, *args):
-        """Warning: this is just empty shell for code implemented in other class."""
-
     @abstractmethod
     def add_progress_bar_metrics(self, *args):
         """Warning: this is just empty shell for code implemented in other class."""
@@ -572,15 +569,15 @@ def process_train_step_outputs(self, all_train_step_outputs, early_stopping_accu
             auto_reduce_tng_result = isinstance(sample_output, Result) and sample_output.should_reduce_on_epoch_end
 
             # only track when a) it needs to be autoreduced OR b) the user wants to manually reduce on epoch end
-            if self.is_overridden('training_epoch_end', model=self.get_model()) or auto_reduce_tng_result:
+            if is_overridden('training_epoch_end', model=self.get_model()) or auto_reduce_tng_result:
                 epoch_end_outputs.append(optimizer_idx_outputs)
 
         return epoch_end_outputs
 
     def check_checkpoint_callback(self, should_check_val):
         # when no val loop is present or fast-dev-run still need to call checkpoints
         # TODO bake this logic into the checkpoint callback
-        should_activate = not self.is_overridden('validation_step') and not should_check_val
+        should_activate = not is_overridden('validation_step', self.get_model()) and not should_check_val
         if should_activate:
             checkpoint_callbacks = [c for c in self.callbacks if isinstance(c, ModelCheckpoint)]
             [c.on_validation_end(self, self.get_model()) for c in checkpoint_callbacks]
@@ -642,7 +639,7 @@ def run_training_epoch_end(self, epoch_output, checkpoint_accumulator, early_sto
         # --------------------------
         # EPOCH END STEP IF DEFINED
         # --------------------------
-        if self.is_overridden('training_epoch_end', model=model):
+        if is_overridden('training_epoch_end', model=model):
             self.global_step += 1
 
             if is_result_obj:

tests/core/test_datamodules.py

@@ -9,6 +9,7 @@
 from tests.base import EvalModelTemplate
 from tests.base.datamodules import TrialMNISTDataModule
 from tests.base.develop_utils import reset_seed
+from pytorch_lightning.utilities.model_utils import is_overridden
 
 
 def test_can_prepare_data(tmpdir):
@@ -348,7 +349,7 @@ def transfer_batch_to_device(self, data, device):
     trainer = Trainer()
     # running .fit() would require us to implement custom data loaders, we mock the model reference instead
     trainer.get_model = MagicMock(return_value=model)
-    if trainer.is_overridden('transfer_batch_to_device', dm):
+    if is_overridden('transfer_batch_to_device', dm):
         model.transfer_batch_to_device = dm.transfer_batch_to_device
 
     batch_gpu = trainer.transfer_batch_to_gpu(batch, 0)