Repository: jdb78/pytorch-forecasting

Issue:
The predict function concatenates the outputs which leads to GPU OOM for large datasets
- PyTorch-Forecasting version: 0.8.2
- PyTorch version: 1.7.1


### Expected behavior
I should be able to loop through a large dataset with `predict` and not run out of memory.

### Actual behavior
If I have a large validation set that I want to go through after the training, then it causes a GPU OOM error, because all the outputs are stored on the GPU and not transferred back to the GPU.

So moving the outputs to CPU or converting them to numpy inside `predict` during the iterations might be a better way, or allow a method that accepts the batches directly (instead of a dataset, dataloader or dataframe).

E.g. somewhere around [here](https://github.com/jdb78/pytorch-forecasting/blob/master/pytorch_forecasting/models/base_model.py#L867) use `out.cpu()` or something similar.

### Code to reproduce the problem
```
# create large dataset
# create data loader (or just use the ds directly)
raw_predictions, x = ptf_module.predict(val_dl, mode="raw", return_x=True)
```


Edit:
Actually `out` is a dictionary, so:
```
for k in out:
    out[k] = out[k].cpu().detach()
```


Discussion:
That is a good point. Will add a device keyword in the predict method.
Or it could also help to have a separate method that accepts the tensors returned by the dataloader iterator (or `dataset.__getitem__`), and not a dataset/dataloader/dataframe.

The usecase is that I want to spot check part of the validation set, and with small datasets it was fine to pass the whole dataset to predict and plot specific indices from the returned predictions, but with a large dataset I run into the above mentioned problem.

A workaround that I'm using at the moment is this:
```
# data is the training dataframe
raw_predictions, x = ptf_module.predict(data[:idx], mode="raw", return_x=True)
```
when the daraframe is converted to the dataset, it only uses the last sample because of the `predict=True` flag, so essentially I get the prediction at `idx`.
Hm. Essentially, you are are asking for a way to filter the dataset. That could be indeed useful. One can already manually modify the index but that is probably a bit cumbersome.

Resolving patch:
diff --git a/pytorch_forecasting/data/timeseries.py b/pytorch_forecasting/data/timeseries.py
--- a/pytorch_forecasting/data/timeseries.py
+++ b/pytorch_forecasting/data/timeseries.py
@@ -4,17 +4,15 @@
 Timeseries data is special and has to be processed and fed to algorithms in a special way. This module
 defines a class that is able to handle a wide variety of timeseries data problems.
 """
-from copy import deepcopy
+from copy import copy as _copy, deepcopy
 from functools import lru_cache
 import inspect
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Callable, Dict, List, Tuple, Union
 import warnings
 
 import matplotlib.pyplot as plt
 import numpy as np
-from numpy.lib.type_check import nan_to_num
 import pandas as pd
-from pandas.core.algorithms import isin
 from sklearn.exceptions import NotFittedError
 from sklearn.preprocessing import RobustScaler, StandardScaler
 from sklearn.utils import shuffle
@@ -1238,6 +1236,73 @@ def _construct_index(self, data: pd.DataFrame, predict_mode: bool) -> pd.DataFra
 
         return df_index
 
+    def filter(self, filter_func: Callable, copy: bool = True) -> "TimeSeriesDataSet":
+        """
+        Filter subsequences in dataset.
+
+        Uses interpretable version of index :py:meth:`~decoded_index`
+        to filter subsequences in dataset.
+
+        Args:
+            filter_func (Callable): function to filter. Should take :py:meth:`~decoded_index`
+                dataframe as only argument which contains group ids and time index columns.
+            copy (bool): if to return copy of dataset or filter inplace.
+
+        Returns:
+            TimeSeriesDataSet: filtered dataset
+        """
+        # calculate filter
+        filtered_index = self.index[np.asarray(filter_func(self.decoded_index))]
+        # raise error if filter removes all entries
+        if len(filtered_index) == 0:
+            raise ValueError("After applying filter no sub-sequences left in dataset")
+        if copy:
+            dataset = _copy(self)
+            dataset.index = filtered_index
+            return dataset
+        else:
+            self.index = filtered_index
+            return self
+
+    @property
+    def decoded_index(self) -> pd.DataFrame:
+        """
+        Get interpretable version of index.
+
+        DataFrame contains
+        - group_id columns in original encoding
+        - time_idx_first column: first time index of subsequence
+        - time_idx_last columns: last time index of subsequence
+        - time_idx_first_prediction columns: first time index which is in decoder
+
+        Returns:
+            pd.DataFrame: index that can be understood in terms of original data
+        """
+        # get dataframe to filter
+        index_start = self.index["index_start"].to_numpy()
+        index_last = self.index["index_end"].to_numpy()
+        index = (
+            # get group ids in order of index
+            pd.DataFrame(self.data["groups"][index_start].numpy(), columns=self.group_ids)
+            # to original values
+            .apply(lambda x: self.transform_values(name=x.name, values=x, group_id=True, inverse=True))
+            # add time index
+            .assign(
+                time_idx_first=self.data["time"][index_start].numpy(),
+                time_idx_last=self.data["time"][index_last].numpy(),
+                # prediction index is last time index - decoder length + 1
+                time_idx_first_prediction=lambda x: x.time_idx_last
+                # decoder length is minimum of
+                - (
+                    x.time_idx_last - x.time_idx_first + 1 - self.min_encoder_length
+                )  # sequence lenght - min decoder length
+                .clip(upper=self.max_prediction_length)  # maximum prediction length
+                .clip(upper=x.time_idx_last - (self.min_prediction_idx - 1))  # not going beyond min prediction idx
+                + 1,
+            )
+        )
+        return index
+
     def plot_randomization(
         self, betas: Tuple[float, float] = None, length: int = None, min_length: int = None
     ) -> Tuple[plt.Figure, torch.Tensor]:
@@ -1337,7 +1402,6 @@ def __getitem__(self, idx: int) -> Tuple[Dict[str, torch.Tensor], torch.Tensor]:
             weight = None
         else:
             weight = self.data["weight"][index.index_start : index.index_end + 1].clone()
-
         # get target scale in the form of a list
         target_scale = self.target_normalizer.get_parameters(groups, self.group_ids)
         if not isinstance(self.target_normalizer, MultiNormalizer):
diff --git a/pytorch_forecasting/models/base_model.py b/pytorch_forecasting/models/base_model.py
--- a/pytorch_forecasting/models/base_model.py
+++ b/pytorch_forecasting/models/base_model.py
@@ -824,6 +824,7 @@ def predict(
 
                 # make prediction
                 out = self(x, **kwargs)  # raw output is dictionary
+
                 out["prediction"] = self.transform_output(out)
 
                 lengths = x["decoder_lengths"]
@@ -864,6 +865,18 @@ def predict(
                 else:
                     raise ValueError(f"Unknown mode {mode} - see docs for valid arguments")
 
+                # move back to CPU
+                if isinstance(out, dict):
+                    for name in out.keys():
+                        if isinstance(out[name], (tuple, list)) and out[name][0].device != torch.device("cpu"):
+                            out[name] = [oi.to("cpu") for oi in out[name]]
+                        elif out[name].device != torch.device("cpu"):
+                            out[name] = out[name].to("cpu")
+                elif isinstance(out, torch.Tensor) and out.device != torch.device("cpu"):
+                    out = out.to("cpu")
+                elif isinstance(out, (list, tuple)) and out[0].device != torch.device("cpu"):
+                    out = [o.to("cpu") for o in out]
+
                 output.append(out)
                 if return_x:
                     x_list.append(x)


Test patch:
diff --git a/tests/test_data.py b/tests/test_data.py
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -565,3 +565,23 @@ def test_lagged_variables(test_data, kwargs):
             target = x[..., target_idx][:, 0]
             lagged_target = torch.roll(x[..., lag_idx], -lag, dims=1)[:, 0]
             assert torch.isclose(target, lagged_target).all(), "lagged target must be the same as non-lagged target"
+
+
+@pytest.mark.parametrize(
+    "agency,first_prediction_idx,should_raise",
+    [("Agency_01", 0, False), ("xxxxx", 0, True), ("Agency_01", 100, True), ("Agency_01", 4, False)],
+)
+def test_filter_data(test_dataset, agency, first_prediction_idx, should_raise):
+    func = lambda x: (x.agency == agency) & (x.time_idx_first_prediction >= first_prediction_idx)
+    if should_raise:
+        with pytest.raises(ValueError):
+            test_dataset.filter(func)
+    else:
+        filtered_dataset = test_dataset.filter(func)
+        assert len(test_dataset.index) > len(
+            filtered_dataset.index
+        ), "filtered dataset should have less entries than original dataset"
+        for x, _ in iter(filtered_dataset.to_dataloader()):
+            index = test_dataset.x_to_index(x)
+            assert (index["agency"] == agency).all(), "Agency filter has failed"
+            assert index["time_idx"].min() == first_prediction_idx, "First prediction filter has failed"