Provide access to unwrapped model in Lite (#12597)

Co-authored-by: Akihiro Nitta <[email protected]>

Author: awaelchli

CHANGELOG.md

@@ -197,10 +197,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue with unsupported torch.inference_mode() on hpu backends by making it use no_grad ([#13014](https://github.com/PyTorchLightning/pytorch-lightning/pull/13014))
 
 
-- Avoid redundant callback restore warning while tuning ([#13026](https://github.com/PyTorchLightning/pytorch-lightning/pull/13026))
+- The model wrapper returned by `LightningLite.setup()` now properly supports pass-through when looking up attributes ([#12597](https://github.com/PyTorchLightning/pytorch-lightning/pull/12597))
 
 
--
+- Avoid redundant callback restore warning while tuning ([#13026](https://github.com/PyTorchLightning/pytorch-lightning/pull/13026))
 
 
 -

pytorch_lightning/lite/lite.py

@@ -152,25 +152,26 @@ def setup(
         *optimizers: Optimizer,
         move_to_device: bool = True,
     ) -> Any:  # no specific return because the way we want our API to look does not play well with mypy
-        """Setup a model and its optimizers for accelerated training.
+        """Set up a model and its optimizers for accelerated training.
 
         Args:
-            model: A model to setup
-            *optimizers: The optimizer(s) to setup (no optimizers is also possible)
+            model: A model to set up
+            *optimizers: The optimizer(s) to set up (no optimizers is also possible)
             move_to_device: If set ``True`` (default), moves the model to the correct device. Set this to ``False``
                 and alternatively use :meth:`to_device` manually.
 
         Returns:
             The tuple of the wrapped model and list of optimizers, in the same order they were passed in.
         """
         self._validate_setup(model, optimizers)
+        original_model = model
 
         if move_to_device:
             model = self._move_model_to_device(model=model, optimizers=list(optimizers))
 
         # Let accelerator/plugin wrap and connect the models and optimizers
         model, optimizers = self._strategy._setup_model_and_optimizers(model, list(optimizers))
-        model = _LiteModule(model, self._precision_plugin)
+        model = _LiteModule(model, self._precision_plugin, original_module=original_model)
         optimizers = [_LiteOptimizer(optimizer=optimizer, strategy=self._strategy) for optimizer in optimizers]
         self._models_setup += 1
         if optimizers:
@@ -181,7 +182,7 @@ def setup(
     def setup_dataloaders(
         self, *dataloaders: DataLoader, replace_sampler: bool = True, move_to_device: bool = True
     ) -> Union[DataLoader, List[DataLoader]]:
-        """Setup one or multiple dataloaders for accelerated training. If you need different settings for each
+        """Set up one or multiple dataloaders for accelerated training. If you need different settings for each
         dataloader, call this method individually for each one.
 
         Args:
@@ -206,7 +207,7 @@ def setup_dataloaders(
     def _setup_dataloader(
         self, dataloader: DataLoader, replace_sampler: bool = True, move_to_device: bool = True
     ) -> DataLoader:
-        """Setup a single dataloader for accelerated training.
+        """Set up a single dataloader for accelerated training.
 
         Args:
             dataloader: The dataloader to accelerate.
@@ -252,10 +253,10 @@ def backward(self, tensor: Tensor, *args: Any, model: Optional[_LiteModule] = No
             **kwargs: Optional named keyword arguments passed to the underlying backward function.
 
         Note:
-            When using ``strategy="deepspeed"`` and multiple models were setup, it is required to pass in the
+            When using ``strategy="deepspeed"`` and multiple models were set up, it is required to pass in the
             model as argument here.
         """
-        module = model.module if model is not None else model
+        module = model._forward_module if model is not None else model
         if isinstance(self._strategy, DeepSpeedStrategy):
             if model is None:
                 if self._models_setup == 0:

pytorch_lightning/lite/wrappers.py

@@ -65,23 +65,29 @@ def step(self, closure: Optional[Callable] = None) -> Any:
 
 
 class _LiteModule(DeviceDtypeModuleMixin):
-    def __init__(self, module: nn.Module, precision_plugin: PrecisionPlugin) -> None:
+    def __init__(
+        self, forward_module: nn.Module, precision_plugin: PrecisionPlugin, original_module: Optional[nn.Module] = None
+    ) -> None:
         """The LiteModule is a thin wrapper around the :class:`torch.nn.Module` and handles precision / autocast
         automatically for the forward pass.
 
         The underlying wrapped module can be accessed via the property :attr:`module`.
 
         Args:
-            module: The module to wrap
+            forward_module: The module to wrap the ``forward`` method on.
             precision_plugin: Reference to the precision plugin for handling precision context
+            original_module: The original, unmodified module as passed into the
+                :meth:`pytorch_lightning.lite.lite.LightningLite.setup` method. This is needed when attribute lookup
+                on this wrapper should pass through to the original module.
         """
         super().__init__()
-        self._module = module
+        self._forward_module = forward_module
+        self._original_module = original_module or forward_module
         self._precision_plugin = precision_plugin
 
     @property
     def module(self) -> nn.Module:
-        return self._module
+        return self._original_module or self._forward_module
 
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Casts all inputs to the right precision and handles autocast for operations in the module forward
@@ -102,12 +108,22 @@ def _convert_float_tensor(t: Tensor) -> Tensor:
         args, kwargs = apply_to_collection([args, kwargs], function=_convert_float_tensor, dtype=Tensor)
 
         with self._precision_plugin.forward_context():
-            output = self.module(*args, **kwargs)
+            output = self._forward_module(*args, **kwargs)
 
         to_type = torch.get_default_dtype()
         output = apply_to_collection(output, function=_convert_float_tensor, dtype=Tensor)
         return output
 
+    def __getattr__(self, item: Any) -> Any:
+        try:
+            # __getattr__ gets called as a last resort if the attribute does not exist
+            # call nn.Module's implementation first
+            return super().__getattr__(item)
+        except AttributeError:
+            # If the attribute is not available on the _LiteModule wrapper, redirect to the wrapped nn.Module
+            original_module = super().__getattr__("_original_module")
+            return getattr(original_module, item)
+
 
 class _LiteDataLoader:
     def __init__(self, dataloader: DataLoader, device: Optional[torch.device] = None) -> None:

tests/lite/test_lite.py

@@ -14,7 +14,7 @@
 import os
 from copy import deepcopy
 from unittest import mock
-from unittest.mock import MagicMock, Mock, PropertyMock
+from unittest.mock import ANY, MagicMock, Mock, PropertyMock
 
 import pytest
 import torch
@@ -80,6 +80,18 @@ def run(self, *args, **kwargs):
     assert lite.run_kwargs == {"three": 3}
 
 
+@mock.patch("pytorch_lightning.strategies.ddp.DistributedDataParallel")
+def test_setup_model(ddp_mock):
+    """Test that the setup method lets the strategy wrap the model, but keeps a reference to the original model."""
+    lite = EmptyLite(accelerator="cpu", strategy="ddp", devices=2)
+    model = nn.Linear(1, 2)
+    lite_model = lite.setup(model)
+    ddp_mock.assert_called_with(module=model, device_ids=ANY)
+    assert lite_model.module == model
+    assert lite_model.weight is model.weight
+    assert lite_model.forward != model.forward
+
+
 def test_setup_optimizers():
     """Test that setup_optimizers can handle no optimizers, one optimizer, or multiple optimizers."""
     lite = EmptyLite()

tests/lite/test_wrappers.py

@@ -33,6 +33,41 @@ def test_lite_module_wraps():
     module = Mock()
     assert _LiteModule(module, Mock()).module is module
 
+    wrapped_module = Mock()
+    original_module = Mock()
+    assert _LiteModule(wrapped_module, Mock(), original_module=original_module).module is original_module
+
+
+def test_lite_module_attribute_lookup():
+    """Test that attribute lookup passes through to the original model when possible."""
+
+    class OriginalModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.layer = torch.nn.Linear(2, 3)
+            self.attribute = 1
+
+        def method(self):
+            return 2
+
+    original_module = OriginalModule()
+
+    class ModuleWrapper(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.wrapped = original_module
+
+    wrapped_module = ModuleWrapper()
+
+    lite_module = _LiteModule(wrapped_module, Mock(), original_module=original_module)
+    assert lite_module.attribute == 1
+    assert lite_module.layer is original_module.layer
+    assert lite_module.method() == 2
+    assert lite_module.forward.__self__.__class__ == _LiteModule
+
+    with pytest.raises(AttributeError):
+        _ = lite_module.not_exists
+
 
 @RunIf(min_gpus=1)
 @pytest.mark.parametrize(