Repository: Lightning-AI/pytorch-lightning

Issue:
Auto move input to proper device for inference
Does PyTorch Lightning provide abstractions for inference? In particular, does it provide ways of automatically handling the transfer to/from GPU when I call `model(x)`, or do I need to roll my own code for that?

## Example Use Case

I have a use case where I train a model on slices of a sliding window of an audio spectrogram (i.e., let's say 1 second chunks). When training is finished, I'd like to see the performance of the model on an entire file. Pseudocode:

```python
# generate training data
X, Y = [], []
for audio_file in audio_files:
    for x, y in sliding_window(audio_file):
        X.append(x); Y.append(y)
X, Y = shuffle(X, Y)  # shuffle the slices of all files

# Train model on slices
model = ExampleModel(X, Y)
trainer = Trainer(gpus=1)
trainer.fit(model)

# Plot the performance on a whole test file:
test_Y = []
for x, _ in sliding_window(test_file)
    test_Y.append(model(x))
plt.plot(test_Y)
```

Notice that during training, the notion of a file is entirely gone, but when I plot my test file, I reintroduce that. Of course, in my real code, my training data `X, Y` is split into training, validation and test, as usual. The plotting step is an additional verification; sort of like putting the pieces together.

## Problem
When the model runs on the GPU, The last part of the code becomes:
```python
# Plot the performance on a whole test file:
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
test_Y = []
for x, _ in sliding_window(test_file)
    y = model(x.to(device)).cpu()
    test_Y.append(y)
plt.plot(test_Y)
```

This isn't the end of the world, but it's not as nice as the other code that PyTorch Lightning helped me refactor. I also can't call `x.type_as(...)` since in that loop, I have no reference type that lives on the CPU/GPU that I could refer to (or maybe I can, but I haven't figured it out).

A workaround to this is to save the model and load it again, on a CPU.

```python
# Train model on slices
# ...
trainer.fit(model)
trainer.save_checkpoint("model.ckpt")
model = ExampleModel.load_from_checkpoint("model.ckpt")

# Plot the performance on a whole test file:
model.eval()
test_Y = []
for x, _ in sliding_window(test_file)
    test_Y.append(model(x))
plt.plot(test_Y)
```

While this removes the noise of the `.to(device)` and `.cpu()` calls, it adds the overhead of having to save the model every time. I also still have to manually call `model.eval()`. The use case of running my model on an entire audio file is not for metrics but for visual inspection; as such I always only sample a few audio files. Running the model on a CPU instead of a GPU for inference thus isn't a problem.

## Question
Is there a more elegant way to achieve the above?


Discussion:
it may be a nice feature to have ddp/tpu also for inference time...
any thoughts @PyTorchLightning/core-contributors @williamFalcon?
We should definitely automate this! 
Let's do it?
@tcwalther pls let me bring some light to it... what are the goal, always use the best you have responsible (GPU/TPU)?

thinking about the inference, the case I see could be to have it as a parameter... because in my case (using notebook) I have GPU but most of the "production" models won't fit there so it will crash
also, you may want to keep the output on GPU so by the inference moving it back to CPU and back to GPU does not make sense

whenever a lightningModule is used 
```
model(x)
```
we put x on the proper device...

@tcwalther this is what you mean no?
> we put x on the proper device...

does it also mean that we shall estimate if the model fits available device?
just thinking about the case with a notebook with low-end GPU so the general request will return GPU is available but it does not reflect that the model does not fit in...

@tcwalther ^^

Resolving patch:
diff --git a/pytorch_lightning/core/decorators.py b/pytorch_lightning/core/decorators.py
--- a/pytorch_lightning/core/decorators.py
+++ b/pytorch_lightning/core/decorators.py
@@ -1,3 +1,9 @@
+from functools import wraps
+from typing import Callable
+
+import torch
+
+from pytorch_lightning.core.lightning import LightningModule
 from pytorch_lightning.utilities import rank_zero_warn
 
 
@@ -12,3 +18,49 @@ def data_loader(fn):
     def inner_fx(self):
         return fn(self)
     return inner_fx
+
+
+def auto_move_data(fn: Callable) -> Callable:
+    """
+    Decorator for :class:`~pytorch_lightning.core.lightning.LightningModule` methods for which
+    input arguments should be moved automatically to the correct device.
+    It as no effect if applied to a method of an object that is not an instance of
+    :class:`~pytorch_lightning.core.lightning.LightningModule` and is typically applied to ``__call__``
+    or ``forward``.
+
+    Args:
+        fn: A LightningModule method for which the arguments should be moved to the device
+            the parameters are on.
+
+    Example:
+
+        .. code-block:: python
+
+            # directly in the source code
+            class LitModel(LightningModule):
+
+                @auto_move_data
+                def forward(self, x):
+                    return x
+
+            # or outside
+            LitModel.forward = auto_move_data(LitModel.forward)
+
+            model = LitModel()
+            model = model.to('cuda')
+            model(torch.zeros(1, 3))
+
+            # input gets moved to device
+            # tensor([[0., 0., 0.]], device='cuda:0')
+
+    """
+    @wraps(fn)
+    def auto_transfer_args(self, *args, **kwargs):
+        if not isinstance(self, LightningModule):
+            return fn(self, *args, **kwargs)
+
+        args = self.transfer_batch_to_device(args, self.device)
+        kwargs = self.transfer_batch_to_device(kwargs, self.device)
+        return fn(self, *args, **kwargs)
+
+    return auto_transfer_args
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -112,6 +112,9 @@ def forward(self, *args, **kwargs):
         This makes it easy to write a complex system for training with the outputs
         you'd want in a prediction setting.
 
+        You may also find the :func:`~pytorch_lightning.core.decorators.auto_move_data` decorator useful
+        when using the module outside Lightning in a production setting.
+
         Args:
             *args: Whatever you decide to pass into the forward method.
             **kwargs: Keyword arguments are also possible.
diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py
--- a/pytorch_lightning/trainer/__init__.py
+++ b/pytorch_lightning/trainer/__init__.py
@@ -101,6 +101,11 @@ def forward(self, x):
     out = pretrained_model(x)
     api_write({'response': out}
 
+
+You may wish to run the model on a variety of devices. Instead of moving the data
+manually to the correct device, decorate the forward method (or any other method you use for inference)
+with :func:`~pytorch_lightning.core.decorators.auto_move_data` and Lightning will take care of the rest.
+
 ------------
 
 Reproducibility


Test patch:
diff --git a/tests/core/test_decorators.py b/tests/core/test_decorators.py
new file mode 100644
--- /dev/null
+++ b/tests/core/test_decorators.py
@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+from tests.base import EvalModelTemplate
+from pytorch_lightning.core.decorators import auto_move_data
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+@pytest.mark.parametrize(['src_device', 'dest_device'], [
+    pytest.param(torch.device('cpu'), torch.device('cpu')),
+    pytest.param(torch.device('cpu', 0), torch.device('cuda', 0)),
+    pytest.param(torch.device('cuda', 0), torch.device('cpu')),
+    pytest.param(torch.device('cuda', 0), torch.device('cuda', 0)),
+])
+def test_auto_move_data(src_device, dest_device):
+    """ Test that the decorator moves the data to the device the model is on. """
+
+    class CurrentModel(EvalModelTemplate):
+        pass
+
+    # apply the decorator
+    CurrentModel.forward = auto_move_data(CurrentModel.forward)
+
+    model = CurrentModel()
+    model = model.to(dest_device)
+    model.prepare_data()
+    loader = model.train_dataloader()
+    x, y, = next(iter(loader))
+    x = x.flatten(1)
+
+    # test that data on source device gets moved to destination device
+    x = x.to(src_device)
+    assert model(x).device == dest_device, "Automoving data to same device as model failed"