sketch

etna/datasets/tsdataset.py

@@ -37,7 +37,7 @@
     from etna.transforms.base import Transform
 
 if SETTINGS.torch_required:
-    from torch.utils.data import Dataset
+    from torch.utils.data import Dataset, IterableDataset
 
 
 class TSDataset:
@@ -1829,6 +1829,17 @@ def info(self, segments: Optional[Sequence[str]] = None) -> None:
         result_string = "\n".join(lines)
         print(result_string)
 
+    def to_torch_generator_dataset(
+        self, make_samples: Callable[[pd.DataFrame], Union[Iterator[dict], Iterable[dict]]], dropna: bool = True
+    ) -> "Dataset":
+        df = self.to_pandas(flatten=True)
+        if dropna:
+            df = df.dropna()  # TODO: Fix this
+
+        # indexes = "generate indexes"
+
+        return _TorchGeneratorDataset(df=df, indexes=indexes, make_samples=make_samples)
+
     def to_torch_dataset(
         self, make_samples: Callable[[pd.DataFrame], Union[Iterator[dict], Iterable[dict]]], dropna: bool = True
     ) -> "Dataset":

etna/datasets/utils.py

@@ -15,7 +15,7 @@
 from etna import SETTINGS
 
 if SETTINGS.torch_required:
-    from torch.utils.data import Dataset
+    from torch.utils.data import Dataset, IterableDataset
 else:
     from unittest.mock import Mock
 
@@ -195,6 +195,24 @@ def duplicate_data(df: pd.DataFrame, segments: Sequence[str], format: str = Data
     return df_long
 
 
+class _TorchGeneratorDataset(Dataset):
+    """In memory dataset for torch dataloader."""
+
+    def __init__(self, df, indexes, make_samples):
+        self.df = df
+        self.indexes = indexes
+        self.make_samples = make_samples
+
+    def __getitem__(self, index):
+        # choose segment dataframe by index
+        # choose slice of segment dataframe by index
+        sample = [sample for sample in self.make_samples(df_segment)]
+        return sample[0]
+
+    def __len__(self):
+        return len(self.indexes)
+
+
 class _TorchDataset(Dataset):
     """In memory dataset for torch dataloader."""
 

etna/models/base.py

@@ -485,6 +485,7 @@ def __init__(
         decoder_length: int,
         train_batch_size: int,
         test_batch_size: int,
+        generate_running_samples: bool,
         trainer_params: Optional[dict],
         train_dataloader_params: Optional[dict],
         test_dataloader_params: Optional[dict],
@@ -527,6 +528,7 @@ def __init__(
         self.decoder_length = decoder_length
         self.train_batch_size = train_batch_size
         self.test_batch_size = test_batch_size
+        self.generate_running_samples = generate_running_samples
         self.train_dataloader_params = {} if train_dataloader_params is None else train_dataloader_params
         self.test_dataloader_params = {} if test_dataloader_params is None else test_dataloader_params
         self.val_dataloader_params = {} if val_dataloader_params is None else val_dataloader_params
@@ -553,12 +555,20 @@ def fit(self, ts: TSDataset) -> "DeepBaseModel":
         :
             Model after fit
         """
-        torch_dataset = ts.to_torch_dataset(
-            functools.partial(
-                self.net.make_samples, encoder_length=self.encoder_length, decoder_length=self.decoder_length
-            ),
-            dropna=True,
-        )
+        if self.generate_running_samples:
+            torch_dataset = ts.to_torch_generator_dataset(
+                functools.partial(
+                    self.net.make_samples, encoder_length=self.encoder_length, decoder_length=self.decoder_length
+                ),
+                dropna=True,
+            )
+        else:
+            torch_dataset = ts.to_torch_dataset(
+                functools.partial(
+                    self.net.make_samples, encoder_length=self.encoder_length, decoder_length=self.decoder_length
+                ),
+                dropna=True,
+            )
         self.raw_fit(torch_dataset)
         return self
 

etna/models/nn/rnn.py

@@ -283,6 +283,7 @@ def __init__(
         loss: Optional["torch.nn.Module"] = None,
         train_batch_size: int = 16,
         test_batch_size: int = 16,
+        generate_running_samples: bool = False,
         optimizer_params: Optional[dict] = None,
         trainer_params: Optional[dict] = None,
         train_dataloader_params: Optional[dict] = None,
@@ -353,6 +354,7 @@ def __init__(
             encoder_length=encoder_length,
             train_batch_size=train_batch_size,
             test_batch_size=test_batch_size,
+            generate_running_samples=generate_running_samples,
             train_dataloader_params=train_dataloader_params,
             test_dataloader_params=test_dataloader_params,
             val_dataloader_params=val_dataloader_params,