wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-24 13:20:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

WIP: causal model (#31)

Browse Source 
* intial causal model

* added training

* initial inference

* formatting

* added readme to model

* key is called control

* added schema of the model

* formatting

* get rid of confusing figure.

* typi

* more typos

* clarification

* added some examples

* fix covariate convention
This commit is contained in:
Kashif Rasul
2021-03-30 16:55:13 +02:00
committed by GitHub Enterprise
parent b01e7c6a24
commit 4e97483fab
4 changed files with 905 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/causal_deepar/README.md
+53
View File
@@ -0,0 +1,53 @@
# Causal `DeepAR`
Causal `DeepAR` model augments the `DeepAR` model by incorporating a causal structure via a `control` time-dependent signal.
The main assumption of this model is that the `target` at time `t` depends not only on the covariates up till time `t` and but also on `control` till time `t`. Thus we encode this structure via a model which now adds a `control_output` distribution layer. The main assumption we have is that at training time our dataset of time series now have an additional `control` key with the corresponding 1-d control-variate values (with the array being as large as `target`).
We translate this structure into the following schematics of Causal-`DeepAR` at time `t`:
```
┌─────────┐ ┌────────┐
│control_t│──────┐ │target_t│
└─────────┘ │ └────────┘
▲ │ ▲
│ │ │
log-prob/ │ log-prob/
sample │ sample
│ ground-truth/ │
┌────────┐ sample/ ┌───────┐
│control │ do │target │
│ dist │ └──────▶│ dist │
└────────┘ └───────┘
▲ ▲
└────────────┬──────────┘
│
┌────┐
───h_t-1─▶│RNN │───h_t──▶
└────┘
▲
┌─────┴────┐
│target_t-1│
└─────┬────┘
┌─────┴─────┐
│control_t-1│
└─────┬─────┘
┌───┴──┐
│cov_t │
└──────┘
```
The model is trained as per the DeepAR assumption, which mainly implies that the covariates and `control` are known for all the time points while training and only the covariates are known for all the time points that we wish to forecast for.
In terms of the `control` at inference time, we now have two choices:
1. We can predict by setting the `control` array to `np.Nan` for the duration of the prediction length in which case the model will sample values from the `control_output.distribution` and feed it auto-regressively to the `target` head and RNN.
2. On the other hand we can in the prediction window "do" an intervention by setting the values of `control` to some fixed value of our choosing for the time steps in the future we are interested in. In this case this model will "break" the causal connections and just use the supplied values and feed those to the `target` head and RNN at the appropriate time points.
## Possible uses
* `target` can be sales and `control` can be discounts
* `target` can be sales and `control` can be the temperature
pts/model/causal_deepar/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .causal_deepar_estimator import CausalDeepAREstimator
from .causal_deepar_network import CausalDeepARNetwork, CausalDeepARTrainingNetwork
pts/model/causal_deepar/causal_deepar_estimator.py
+280
View File
@@ -0,0 +1,280 @@
from typing import List, Optional
import numpy as np
import torch
import torch.nn as nn
from gluonts.core.component import validated
from gluonts.dataset.field_names import FieldName
from gluonts.time_feature import (
TimeFeature,
get_lags_for_frequency,
time_features_from_frequency_str,
)
from gluonts.transform import (
Transformation,
Chain,
RemoveFields,
SetField,
AsNumpyArray,
AddObservedValuesIndicator,
AddTimeFeatures,
AddAgeFeature,
VstackFeatures,
InstanceSplitter,
ValidationSplitSampler,
TestSplitSampler,
ExpectedNumInstanceSampler,
)
from gluonts.torch.support.util import copy_parameters
from gluonts.torch.model.predictor import PyTorchPredictor
from gluonts.torch.modules.distribution_output import DistributionOutput
from gluonts.model.predictor import Predictor
from pts.model.utils import get_module_forward_input_names
from pts import Trainer
from pts.model import PyTorchEstimator
from pts.modules import StudentTOutput
from .causal_deepar_network import (
CausalDeepARTrainingNetwork,
CausalDeepARPredictionNetwork,
)
class CausalDeepAREstimator(PyTorchEstimator):
@validated()
def __init__(
self,
freq: str,
prediction_length: int,
input_size: int,
trainer: Trainer = Trainer(),
context_length: Optional[int] = None,
num_layers: int = 2,
num_cells: int = 40,
cell_type: str = "LSTM",
dropout_rate: float = 0.1,
use_feat_dynamic_real: bool = False,
use_feat_dynamic_cat: bool = False,
use_feat_static_cat: bool = False,
use_feat_static_real: bool = False,
cardinality: Optional[List[int]] = None,
embedding_dimension: Optional[List[int]] = None,
distr_output: DistributionOutput = StudentTOutput(),
control_output: DistributionOutput = StudentTOutput(),
scaling: bool = True,
lags_seq: Optional[List[int]] = None,
time_features: Optional[List[TimeFeature]] = None,
num_parallel_samples: int = 100,
dtype: np.dtype = np.float32,
) -> None:
super().__init__(trainer=trainer)
self.freq = freq
self.context_length = (
context_length if context_length is not None else prediction_length
)
self.prediction_length = prediction_length
self.distr_output = distr_output
self.control_output = control_output
self.distr_output.dtype = dtype
self.input_size = input_size
self.num_layers = num_layers
self.num_cells = num_cells
self.cell_type = cell_type
self.dropout_rate = dropout_rate
self.use_feat_dynamic_real = use_feat_dynamic_real
self.use_feat_dynamic_cat = use_feat_dynamic_cat
self.use_feat_static_cat = use_feat_static_cat
self.use_feat_static_real = use_feat_static_real
self.cardinality = cardinality if cardinality and use_feat_static_cat else [1]
self.embedding_dimension = (
embedding_dimension
if embedding_dimension is not None
else [min(50, (cat + 1) // 2) for cat in self.cardinality]
)
self.scaling = scaling
self.lags_seq = (
lags_seq
if lags_seq is not None
else get_lags_for_frequency(freq_str=freq, lag_ub=self.context_length)
)
self.time_features = (
time_features
if time_features is not None
else time_features_from_frequency_str(self.freq)
)
self.history_length = self.context_length + max(self.lags_seq)
self.num_parallel_samples = num_parallel_samples
self.train_sampler = ExpectedNumInstanceSampler(
num_instances=1.0, min_future=prediction_length
)
self.validation_sampler = ValidationSplitSampler(min_future=prediction_length)
def create_transformation(self) -> Transformation:
remove_field_names = []
if not self.use_feat_static_real:
remove_field_names.append(FieldName.FEAT_STATIC_REAL)
if not self.use_feat_dynamic_real:
remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)
if not self.use_feat_dynamic_cat:
remove_field_names.append(FieldName.FEAT_DYNAMIC_CAT)
return Chain(
[RemoveFields(field_names=remove_field_names)]
+ (
[SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0])]
if not self.use_feat_static_cat
else []
)
+ (
[SetField(output_field=FieldName.FEAT_STATIC_REAL, value=[0.0])]
if not self.use_feat_static_real
else []
)
+ [
AsNumpyArray(
field=FieldName.FEAT_STATIC_CAT,
expected_ndim=1,
dtype=np.long,
),
AsNumpyArray(
field=FieldName.FEAT_STATIC_REAL,
expected_ndim=1,
dtype=self.dtype,
),
AsNumpyArray(
field=FieldName.TARGET,
# in the following line, we add 1 for the time dimension
expected_ndim=1 + len(self.distr_output.event_shape),
dtype=self.dtype,
),
AsNumpyArray(
field="control",
# in the following line, we add 1 for the time dimension
expected_ndim=1 + len(self.control_output.event_shape),
dtype=self.dtype,
),
AddObservedValuesIndicator(
target_field=FieldName.TARGET,
output_field=FieldName.OBSERVED_VALUES,
dtype=self.dtype,
),
AddTimeFeatures(
start_field=FieldName.START,
target_field=FieldName.TARGET,
output_field=FieldName.FEAT_TIME,
time_features=self.time_features,
pred_length=self.prediction_length,
),
AddAgeFeature(
target_field=FieldName.TARGET,
output_field=FieldName.FEAT_AGE,
pred_length=self.prediction_length,
log_scale=True,
dtype=self.dtype,
),
VstackFeatures(
output_field=FieldName.FEAT_TIME,
input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
+ (
[FieldName.FEAT_DYNAMIC_REAL]
if self.use_feat_dynamic_real
else []
)
+ (
[FieldName.FEAT_DYNAMIC_CAT]
if self.use_feat_dynamic_cat
else []
),
),
]
)
def create_instance_splitter(self, mode: str):
assert mode in ["training", "validation", "test"]
instance_sampler = {
"training": self.train_sampler,
"validation": self.validation_sampler,
"test": TestSplitSampler(),
}[mode]
return InstanceSplitter(
target_field=FieldName.TARGET,
is_pad_field=FieldName.IS_PAD,
start_field=FieldName.START,
forecast_start_field=FieldName.FORECAST_START,
instance_sampler=instance_sampler,
past_length=self.history_length,
future_length=self.prediction_length,
time_series_fields=[
FieldName.FEAT_TIME,
FieldName.OBSERVED_VALUES,
"control",
],
)
def create_training_network(
self, device: torch.device
) -> CausalDeepARTrainingNetwork:
return CausalDeepARTrainingNetwork(
input_size=self.input_size,
num_layers=self.num_layers,
num_cells=self.num_cells,
cell_type=self.cell_type,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
distr_output=self.distr_output,
control_output=self.control_output,
dropout_rate=self.dropout_rate,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=self.scaling,
dtype=self.dtype,
).to(device)
def create_predictor(
self,
transformation: Transformation,
trained_network: nn.Module,
device: torch.device,
) -> Predictor:
prediction_network = CausalDeepARPredictionNetwork(
num_parallel_samples=self.num_parallel_samples,
input_size=self.input_size,
num_layers=self.num_layers,
num_cells=self.num_cells,
cell_type=self.cell_type,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
distr_output=self.distr_output,
control_output=self.control_output,
dropout_rate=self.dropout_rate,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=self.scaling,
dtype=self.dtype,
).to(device)
copy_parameters(trained_network, prediction_network)
input_names = get_module_forward_input_names(prediction_network)
prediction_splitter = self.create_instance_splitter("test")
return PyTorchPredictor(
input_transform=transformation + prediction_splitter,
input_names=input_names,
prediction_net=prediction_network,
batch_size=self.trainer.batch_size,
freq=self.freq,
prediction_length=self.prediction_length,
device=device,
)
pts/model/causal_deepar/causal_deepar_network.py
+570
View File
@@ -0,0 +1,570 @@
from typing import List, Optional, Tuple, Union
import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Distribution
from gluonts.core.component import validated
from gluonts.torch.modules.distribution_output import DistributionOutput
from pts.model import weighted_average
from pts.modules import MeanScaler, NOPScaler, FeatureEmbedder
def prod(xs):
p = 1
for x in xs:
p *= x
return p
class CausalDeepARNetwork(nn.Module):
@validated()
def __init__(
self,
input_size: int,
num_layers: int,
num_cells: int,
cell_type: str,
history_length: int,
context_length: int,
prediction_length: int,
distr_output: DistributionOutput,
control_output: DistributionOutput,
dropout_rate: float,
cardinality: List[int],
embedding_dimension: List[int],
lags_seq: List[int],
scaling: bool = True,
dtype: np.dtype = np.float32,
) -> None:
super().__init__()
self.num_layers = num_layers
self.num_cells = num_cells
self.cell_type = cell_type
self.history_length = history_length
self.context_length = context_length
self.prediction_length = prediction_length
self.dropout_rate = dropout_rate
self.cardinality = cardinality
self.embedding_dimension = embedding_dimension
self.num_cat = len(cardinality)
self.scaling = scaling
self.dtype = dtype
self.lags_seq = lags_seq
self.distr_output = distr_output
self.control_output = control_output
rnn = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type]
self.rnn = rnn(
input_size=input_size,
hidden_size=num_cells,
num_layers=num_layers,
dropout=dropout_rate,
batch_first=True,
)
self.target_shape = distr_output.event_shape
self.proj_distr_args = distr_output.get_args_proj(num_cells + 1)
self.proj_control_args = control_output.get_args_proj(num_cells)
self.embedder = FeatureEmbedder(
cardinalities=cardinality, embedding_dims=embedding_dimension
)
if scaling:
self.scaler = MeanScaler(keepdim=True)
self.control_scaler = NOPScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
self.control_scaler = NOPScaler(keepdim=True)
@staticmethod
def get_lagged_subsequences(
sequence: torch.Tensor,
sequence_length: int,
indices: List[int],
subsequences_length: int = 1,
) -> torch.Tensor:
"""
Returns lagged subsequences of a given sequence.
Parameters
----------
sequence : Tensor
the sequence from which lagged subsequences should be extracted.
Shape: (N, T, C).
sequence_length : int
length of sequence in the T (time) dimension (axis = 1).
indices : List[int]
list of lag indices to be used.
subsequences_length : int
length of the subsequences to be extracted.
Returns
--------
lagged : Tensor
a tensor of shape (N, S, C, I), where S = subsequences_length and
I = len(indices), containing lagged subsequences. Specifically,
lagged[i, j, :, k] = sequence[i, -indices[k]-S+j, :].
"""
assert max(indices) + subsequences_length <= sequence_length, (
f"lags cannot go further than history length, found lag {max(indices)} "
f"while history length is only {sequence_length}"
)
assert all(lag_index >= 0 for lag_index in indices)
lagged_values = []
for lag_index in indices:
begin_index = -lag_index - subsequences_length
end_index = -lag_index if lag_index > 0 else None
lagged_values.append(sequence[:, begin_index:end_index, ...])
return torch.stack(lagged_values, dim=-1)
def unroll_encoder(
self,
feat_static_cat: torch.Tensor, # (batch_size, num_features)
feat_static_real: torch.Tensor, # (batch_size, num_features)
past_time_feat: torch.Tensor, # (batch_size, history_length, num_features)
past_target: torch.Tensor, # (batch_size, history_length, *target_shape)
past_observed_values: torch.Tensor, # (batch_size, history_length, *target_shape)
past_control: torch.Tensor,
future_control: Optional[torch.Tensor] = None,
future_time_feat: Optional[
torch.Tensor
] = None, # (batch_size, prediction_length, num_features)
future_target: Optional[
torch.Tensor
] = None, # (batch_size, prediction_length, *target_shape)
) -> Tuple[torch.Tensor, Union[torch.Tensor, List], torch.Tensor, torch.Tensor]:
if future_time_feat is None or future_target is None or future_control is None:
time_feat = past_time_feat[
:, self.history_length - self.context_length :, ...
]
sequence = past_target
control_sequence = past_control
sequence_length = self.history_length
subsequences_length = self.context_length
else:
time_feat = torch.cat(
(
past_time_feat[:, self.history_length - self.context_length :, ...],
future_time_feat,
),
dim=1,
)
sequence = torch.cat((past_target, future_target), dim=1)
control_sequence = torch.cat((past_control, future_control), dim=1)
sequence_length = self.history_length + self.prediction_length
subsequences_length = self.context_length + self.prediction_length
lags = self.get_lagged_subsequences(
sequence=sequence,
sequence_length=sequence_length,
indices=self.lags_seq,
subsequences_length=subsequences_length,
)
control_lags = self.get_lagged_subsequences(
sequence=control_sequence,
sequence_length=sequence_length,
indices=self.lags_seq,
subsequences_length=subsequences_length,
)
# scale is computed on the context length last units of the past target
# scale shape is (batch_size, 1, *target_shape)
_, scale = self.scaler(
past_target[:, self.context_length :, ...],
past_observed_values[:, self.context_length :, ...],
)
_, control_scale = self.control_scaler(
past_control[:, self.context_length :, ...],
past_observed_values[:, self.context_length :, ...],
)
# (batch_size, num_features)
embedded_cat = self.embedder(feat_static_cat)
# in addition to embedding features, use the log scale as it can help
# prediction too
# (batch_size, num_features + prod(target_shape))
static_feat = torch.cat(
(
embedded_cat,
feat_static_real,
scale.log() if len(self.target_shape) == 0 else scale.squeeze(1).log(),
control_scale.log()
if len(self.target_shape) == 0
else control_scale.squeeze(1).log(),
),
dim=1,
)
# (batch_size, subsequences_length, num_features + 1)
repeated_static_feat = static_feat.unsqueeze(1).expand(
-1, subsequences_length, -1
)
# (batch_size, sub_seq_len, *target_shape, num_lags)
lags_scaled = lags / scale.unsqueeze(-1)
control_lags_scaled = control_lags / control_scale.unsqueeze(-1)
# from (batch_size, sub_seq_len, *target_shape, num_lags)
# to (batch_size, sub_seq_len, prod(target_shape) * num_lags)
input_lags = lags_scaled.reshape(
(-1, subsequences_length, len(self.lags_seq) * prod(self.target_shape))
)
input_control_lags = control_lags_scaled.reshape(
(-1, subsequences_length, len(self.lags_seq) * prod(self.target_shape))
)
# (batch_size, sub_seq_len, input_dim)
inputs = torch.cat(
(input_lags, input_control_lags, time_feat, repeated_static_feat), dim=-1
)
# unroll encoder
outputs, state = self.rnn(inputs)
# outputs: (batch_size, seq_len, num_cells)
# state: list of (num_layers, batch_size, num_cells) tensors
# scale: (batch_size, 1, *target_shape)
# static_feat: (batch_size, num_features + prod(target_shape))
return outputs, state, scale, control_scale, static_feat
class CausalDeepARTrainingNetwork(CausalDeepARNetwork):
def distribution(
self,
feat_static_cat: torch.Tensor,
feat_static_real: torch.Tensor,
past_time_feat: torch.Tensor,
past_target: torch.Tensor,
past_observed_values: torch.Tensor,
future_time_feat: torch.Tensor,
future_target: torch.Tensor,
future_observed_values: torch.Tensor,
past_control: torch.Tensor,
future_control: torch.Tensor,
) -> Tuple[Distribution]:
rnn_outputs, _, scale, control_scale, _ = self.unroll_encoder(
feat_static_cat=feat_static_cat,
feat_static_real=feat_static_real,
past_time_feat=past_time_feat,
past_target=past_target,
past_observed_values=past_observed_values,
future_time_feat=future_time_feat,
future_target=future_target,
past_control=past_control,
future_control=future_control,
)
control_dist_args = self.proj_control_args(rnn_outputs)
control = torch.cat(
(
past_control[:, self.history_length - self.context_length :, ...],
future_control,
),
dim=1,
)
distr_args = self.proj_distr_args(
torch.cat((rnn_outputs, control.unsqueeze(-1)), dim=-1)
)
return (
self.control_output.distribution(control_dist_args, scale=control_scale),
self.distr_output.distribution(distr_args, scale=scale),
)
def forward(
self,
feat_static_cat: torch.Tensor,
feat_static_real: torch.Tensor,
past_time_feat: torch.Tensor,
past_target: torch.Tensor,
past_observed_values: torch.Tensor,
future_time_feat: torch.Tensor,
future_target: torch.Tensor,
future_observed_values: torch.Tensor,
past_control: torch.Tensor,
future_control: torch.Tensor,
) -> torch.Tensor:
control_distr, distr = self.distribution(
feat_static_cat=feat_static_cat,
feat_static_real=feat_static_real,
past_time_feat=past_time_feat,
past_target=past_target,
past_observed_values=past_observed_values,
future_time_feat=future_time_feat,
future_target=future_target,
future_observed_values=future_observed_values,
past_control=past_control,
future_control=future_control,
)
# put together target sequence
# (batch_size, seq_len, *target_shape)
target = torch.cat(
(
past_target[:, self.history_length - self.context_length :, ...],
future_target,
),
dim=1,
)
control = torch.cat(
(
past_control[:, self.history_length - self.context_length :, ...],
future_control,
),
dim=1,
)
# (batch_size, seq_len)
loss = -control_distr.log_prob(control) - distr.log_prob(target)
# (batch_size, seq_len, *target_shape)
observed_values = torch.cat(
(
past_observed_values[
:, self.history_length - self.context_length :, ...
],
future_observed_values,
),
dim=1,
)
# mask the loss at one time step iff one or more observations is missing in the target dimensions
# (batch_size, seq_len)
loss_weights = (
observed_values
if (len(self.target_shape) == 0)
else observed_values.min(dim=-1, keepdim=False)
)
weighted_loss = weighted_average(loss, weights=loss_weights)
return weighted_loss, loss
class CausalDeepARPredictionNetwork(CausalDeepARNetwork):
def __init__(self, num_parallel_samples: int = 100, **kwargs) -> None:
super().__init__(**kwargs)
self.num_parallel_samples = num_parallel_samples
# for decoding the lags are shifted by one, at the first time-step
# of the decoder a lag of one corresponds to the last target value
self.shifted_lags = [l - 1 for l in self.lags_seq]
def sampling_decoder(
self,
static_feat: torch.Tensor,
past_target: torch.Tensor,
past_control: torch.Tensor,
time_feat: torch.Tensor,
scale: torch.Tensor,
control_scale: torch.Tensor,
future_control: torch.Tensor,
begin_states: Union[torch.Tensor, List[torch.Tensor]],
) -> torch.Tensor:
"""
Computes sample paths by unrolling the RNN starting with a initial
input and state.
Parameters
----------
static_feat : Tensor
static features. Shape: (batch_size, num_static_features).
past_target : Tensor
target history. Shape: (batch_size, history_length).
time_feat : Tensor
time features. Shape: (batch_size, prediction_length, num_time_features).
scale : Tensor
tensor containing the scale of each element in the batch. Shape: (batch_size, 1, 1).
begin_states : List or Tensor
list of initial states for the LSTM layers or tensor for GRU.
the shape of each tensor of the list should be (num_layers, batch_size, num_cells)
Returns
--------
Tensor
A tensor containing sampled paths.
Shape: (batch_size, num_sample_paths, prediction_length).
"""
# blows-up the dimension of each tensor to batch_size * self.num_parallel_samples for increasing parallelism
repeated_past_target = past_target.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
repeated_time_feat = time_feat.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
repeated_past_control = past_control.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
repeated_static_feat = static_feat.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
).unsqueeze(1)
repeated_future_control = future_control.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
repeated_scale = scale.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
repeated_control_scale = control_scale.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
if self.cell_type == "LSTM":
repeated_states = [
s.repeat_interleave(repeats=self.num_parallel_samples, dim=1)
for s in begin_states
]
else:
repeated_states = begin_states.repeat_interleave(
repeats=self.num_parallel_samples, dim=1
)
future_samples = []
# for each future time-units we draw new samples for this time-unit and update the state
for k in range(self.prediction_length):
# (batch_size * num_samples, 1, *target_shape, num_lags)
lags = self.get_lagged_subsequences(
sequence=repeated_past_target,
sequence_length=self.history_length + k,
indices=self.shifted_lags,
subsequences_length=1,
)
control_lags = self.get_lagged_subsequences(
sequence=repeated_past_control,
sequence_length=self.history_length + k,
indices=self.shifted_lags,
subsequences_length=1,
)
# (batch_size * num_samples, 1, *target_shape, num_lags)
lags_scaled = lags / repeated_scale.unsqueeze(-1)
control_lags_scaled = control_lags / repeated_control_scale.unsqueeze(-1)
# from (batch_size * num_samples, 1, *target_shape, num_lags)
# to (batch_size * num_samples, 1, prod(target_shape) * num_lags)
input_lags = lags_scaled.reshape(
(-1, 1, prod(self.target_shape) * len(self.lags_seq))
)
input_control_lags = control_lags_scaled.reshape(
(-1, 1, len(self.lags_seq) * prod(self.target_shape))
)
# (batch_size * num_samples, 1, prod(target_shape) * num_lags + num_time_features + num_static_features)
decoder_input = torch.cat(
(
input_lags,
input_control_lags,
repeated_time_feat[:, k : k + 1, :],
repeated_static_feat,
),
dim=-1,
)
# output shape: (batch_size * num_samples, 1, num_cells)
# state shape: (batch_size * num_samples, num_cells)
rnn_outputs, repeated_states = self.rnn(decoder_input, repeated_states)
control_dist_args = self.proj_control_args(rnn_outputs)
control_distr = self.control_output.distribution(
control_dist_args, repeated_control_scale
)
new_control_sample = control_distr.sample()
control = repeated_future_control[:, k : k + 1]
control[control != control] = new_control_sample[control != control]
repeated_past_control = torch.cat((repeated_past_control, control), dim=1)
distr_args = self.proj_distr_args(
torch.cat((rnn_outputs, control.unsqueeze(-1)), dim=-1)
)
# compute likelihood of target given the predicted parameters
distr = self.distr_output.distribution(distr_args, scale=repeated_scale)
# (batch_size * num_samples, 1, *target_shape)
new_samples = distr.sample()
# (batch_size * num_samples, seq_len, *target_shape)
repeated_past_target = torch.cat((repeated_past_target, new_samples), dim=1)
future_samples.append(new_samples)
# (batch_size * num_samples, prediction_length, *target_shape)
samples = torch.cat(future_samples, dim=1)
# (batch_size, num_samples, prediction_length, *target_shape)
return samples.reshape(
(
(-1, self.num_parallel_samples)
+ (self.prediction_length,)
+ self.target_shape
)
)
# noinspection PyMethodOverriding,PyPep8Naming
def forward(
self,
feat_static_cat: torch.Tensor, # (batch_size, num_features)
feat_static_real: torch.Tensor, # (batch_size, num_features)
past_time_feat: torch.Tensor, # (batch_size, history_length, num_features)
past_target: torch.Tensor, # (batch_size, history_length, *target_shape)
past_observed_values: torch.Tensor, # (batch_size, history_length, *target_shape)
future_time_feat: torch.Tensor, # (batch_size, prediction_length, num_features)
past_control: torch.Tensor,
future_control: torch.Tensor,
) -> torch.Tensor:
"""
Predicts samples, all tensors should have NTC layout.
Parameters
----------
feat_static_cat : (batch_size, num_features)
feat_static_real : (batch_size, num_features)
past_time_feat : (batch_size, history_length, num_features)
past_target : (batch_size, history_length, *target_shape)
past_observed_values : (batch_size, history_length, *target_shape)
future_time_feat : (batch_size, prediction_length, num_features)
Returns
-------
Tensor
Predicted samples
"""
# unroll the decoder in "prediction mode", i.e. with past data only
_, state, scale, control_scale, static_feat = self.unroll_encoder(
feat_static_cat=feat_static_cat,
feat_static_real=feat_static_real,
past_time_feat=past_time_feat,
past_target=past_target,
past_observed_values=past_observed_values,
past_control=past_control,
future_time_feat=None,
future_target=None,
future_control=None,
)
return self.sampling_decoder(
past_target=past_target,
past_control=past_control,
time_feat=future_time_feat,
static_feat=static_feat,
future_control=future_control,
scale=scale,
control_scale=control_scale,
begin_states=state,
)