diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b77f3a6a..44402f9ed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,7 +40,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fix working with `embedding_sizes` in `202-NN_examples` notebook ([#489](https://github.com/etna-team/etna/pull/489)) - Disallow dropping target in `TSDataset.drop_features` ([#491](https://github.com/etna-team/etna/pull/491)) - Optimize memory usage in `TFTNativeModel` by eliminating copying during making samples ([#494](https://github.com/etna-team/etna/pull/494)) -- +- Optimize memory usage in `DeepStateModel` and `DeepARNativeModel` by eliminating copying during making samples ([#499](https://github.com/etna-team/etna/pull/499)) - - - diff --git a/etna/models/nn/deepar_native/deepar.py b/etna/models/nn/deepar_native/deepar.py index ece67ce34..42475c37f 100644 --- a/etna/models/nn/deepar_native/deepar.py +++ b/etna/models/nn/deepar_native/deepar.py @@ -132,6 +132,11 @@ def forward(self, x: DeepARNativeBatch, *args, **kwargs): # type: ignore decoder_categorical = x["decoder_categorical"] # each (batch_size, decoder_length, 1) decoder_target = x["decoder_target"].float() # (batch_size, decoder_length, 1) decoder_length = decoder_real.shape[1] + weights = x["weight"] + + # scale target values at index 0 + encoder_real[:, :, 0] = encoder_real[:, :, 0] / weights.unsqueeze(1) + decoder_real[:, :, 0] = decoder_real[:, :, 0] / weights.unsqueeze(1) encoder_embeddings = self.embedding(encoder_categorical) if self.embedding is not None else torch.Tensor() decoder_embeddings = self.embedding(decoder_categorical) if self.embedding is not None else torch.Tensor() @@ -191,7 +196,11 @@ def step(self, batch: DeepARNativeBatch, *args, **kwargs): # type: ignore decoder_categorical = batch["decoder_categorical"] # each (batch_size, decoder_length, 1) encoder_target = batch["encoder_target"].float() # (batch_size, encoder_length-1, 1) decoder_target = batch["decoder_target"].float() # (batch_size, decoder_length, 1) - weights = batch["weight"] + weights = batch["weight"] # (batch_size) + + # scale target values at index 0 + encoder_real[:, :, 0] = encoder_real[:, :, 0] / weights.unsqueeze(1) + decoder_real[:, :, 0] = decoder_real[:, :, 0] / weights.unsqueeze(1) encoder_embeddings = self.embedding(encoder_categorical) if self.embedding is not None else torch.Tensor() decoder_embeddings = self.embedding(decoder_categorical) if self.embedding is not None else torch.Tensor() @@ -255,11 +264,10 @@ def _make( return None # Get shifted target and concatenate it with real values features - sample["decoder_real"] = values_real[start_idx + encoder_length : start_idx + total_sample_length].copy() + sample["decoder_real"] = values_real[start_idx + encoder_length : start_idx + total_sample_length] # Get shifted target and concatenate it with real values features - sample["encoder_real"] = values_real[start_idx : start_idx + encoder_length].copy() - sample["encoder_real"] = sample["encoder_real"][1:] + sample["encoder_real"] = values_real[start_idx + 1 : start_idx + encoder_length] for index, feature in enumerate(self.embedding_sizes.keys()): sample["encoder_categorical"][feature] = values_categorical[index][ @@ -276,10 +284,6 @@ def _make( sample["segment"] = segment sample["weight"] = 1 + sample["encoder_target"].mean() if self.scale else 1 - sample["encoder_real"][:, 0] = values_real[start_idx + 1 : start_idx + encoder_length, 0] / sample["weight"] - sample["decoder_real"][:, 0] = ( - values_real[start_idx + encoder_length : start_idx + total_sample_length, 0] / sample["weight"] - ) return sample diff --git a/etna/models/nn/deepstate/deepstate.py b/etna/models/nn/deepstate/deepstate.py index d2f72d02a..7250bde03 100644 --- a/etna/models/nn/deepstate/deepstate.py +++ b/etna/models/nn/deepstate/deepstate.py @@ -4,6 +4,7 @@ from typing import Optional from typing import Tuple +import numpy as np import pandas as pd from typing_extensions import TypedDict @@ -117,9 +118,9 @@ def step(self, batch: DeepStateBatch, *args, **kwargs): # type: ignore : loss, true_target, prediction_target """ - encoder_real = batch["encoder_real"] # (batch_size, seq_length, input_size) + encoder_real = batch["encoder_real"].float() # (batch_size, seq_length, input_size) encoder_categorical = batch["encoder_categorical"] # each (batch_size, seq_length, 1) - targets = batch["encoder_target"] # (batch_size, seq_length, 1) + targets = batch["encoder_target"].float() # (batch_size, seq_length, 1) seq_length = targets.shape[1] datetime_index = batch["datetime_index"].permute(1, 0, 2)[ :, :, :seq_length @@ -159,11 +160,11 @@ def forward(self, x: DeepStateBatch, *args, **kwargs): # type: ignore : forecast with shape (batch_size, decoder_length, 1) """ - encoder_real = x["encoder_real"] # (batch_size, seq_length, input_size) + encoder_real = x["encoder_real"].float() # (batch_size, seq_length, input_size) encoder_categorical = x["encoder_categorical"] # each (batch_size, seq_length, 1) seq_length = encoder_real.shape[1] - targets = x["encoder_target"][:, :seq_length] # (batch_size, seq_length, 1) - decoder_real = x["decoder_real"] # (batch_size, horizon, input_size) + targets = x["encoder_target"][:, :seq_length].float() # (batch_size, seq_length, 1) + decoder_real = x["decoder_real"].float() # (batch_size, horizon, input_size) decoder_categorical = x["decoder_categorical"] # each (batch_size, horizon, 1) datetime_index_train = x["datetime_index"].permute(1, 0, 2)[ :, :, :seq_length @@ -213,26 +214,23 @@ def forward(self, x: DeepStateBatch, *args, **kwargs): # type: ignore def make_samples(self, df: pd.DataFrame, encoder_length: int, decoder_length: int) -> Iterator[dict]: """Make samples from segment DataFrame.""" values_real = df.drop(columns=["target", "segment", "timestamp"] + list(self.embedding_sizes.keys())).values - values_real = torch.from_numpy(values_real).float() # Categories that were not seen during `fit` will be filled with new category for feature in self.embedding_sizes: df[feature] = df[feature].astype(float).fillna(self.embedding_sizes[feature][0]) # Columns in `values_categorical` are in the same order as in `embedding_sizes` - values_categorical = torch.from_numpy(df[self.embedding_sizes.keys()].values.T) + values_categorical = df[self.embedding_sizes.keys()].values.T - values_datetime = torch.from_numpy(self.ssm.generate_datetime_index(df["timestamp"])) - values_datetime = values_datetime.to(torch.int64) + values_datetime = self.ssm.generate_datetime_index(df["timestamp"]).astype(int) values_target = df["target"].values - values_target = torch.from_numpy(values_target).float() segment = df["segment"].values[0] def _make( - values_target: torch.Tensor, - values_real: torch.Tensor, - values_categorical: torch.Tensor, - values_datetime: torch.Tensor, + values_target: np.ndarray, + values_real: np.ndarray, + values_categorical: np.ndarray, + values_datetime: np.ndarray, segment: str, start_idx: int, encoder_length: int, diff --git a/tests/test_models/test_nn/deepar_native/test_deepar_native.py b/tests/test_models/test_nn/deepar_native/test_deepar_native.py index 16e009dee..3cb298b8a 100644 --- a/tests/test_models/test_nn/deepar_native/test_deepar_native.py +++ b/tests/test_models/test_nn/deepar_native/test_deepar_native.py @@ -105,12 +105,11 @@ def test_deepar_make_samples(df_name, scale, weights, cat_columns, request): num_samples_check = 2 df["target_shifted"] = df["target"].shift(1) for i in range(num_samples_check): - df[f"target_shifted_scaled_{i}"] = df["target_shifted"] / weights[i] expected_sample = { - "encoder_real": df[[f"target_shifted_scaled_{i}", "regressor_float", "regressor_int"]] + "encoder_real": df[["target_shifted", "regressor_float", "regressor_int"]] .iloc[1 + i : encoder_length + i] .values, - "decoder_real": df[[f"target_shifted_scaled_{i}", "regressor_float", "regressor_int"]] + "decoder_real": df[["target_shifted", "regressor_float", "regressor_int"]] .iloc[encoder_length + i : encoder_length + decoder_length + i] .values, "encoder_categorical": { @@ -138,6 +137,12 @@ def test_deepar_make_samples(df_name, scale, weights, cat_columns, request): assert ts_samples[i]["segment"] == "segment_1" for key in expected_sample: np.testing.assert_equal(ts_samples[i][key], expected_sample[key]) + if "categorical" in key: + for column in ts_samples[i][key]: + assert ts_samples[i][key][column].base is not None + else: + if key != "weight": + assert ts_samples[i][key].base is not None @pytest.mark.parametrize("encoder_length", [1, 2, 10]) diff --git a/tests/test_models/test_nn/test_deepstate.py b/tests/test_models/test_nn/test_deepstate.py index a56b2cd4a..78efb4370 100644 --- a/tests/test_models/test_nn/test_deepstate.py +++ b/tests/test_models/test_nn/test_deepstate.py @@ -139,6 +139,11 @@ def test_deepstate_make_samples(df_name, cat_columns, request): assert ts_samples[i]["segment"] == "segment_1" for key in expected_sample: np.testing.assert_equal(ts_samples[i][key], expected_sample[key]) + if "categorical" in key: + for column in ts_samples[i][key]: + assert ts_samples[i][key][column].base is not None + else: + assert ts_samples[i][key].base is not None def test_save_load(example_tsds):