diff --git a/pts/model/deepvar/deepvar_estimator.py b/pts/model/deepvar/deepvar_estimator.py index dd2e453..f3d3653 100644 --- a/pts/model/deepvar/deepvar_estimator.py +++ b/pts/model/deepvar/deepvar_estimator.py @@ -37,6 +37,7 @@ from .deepvar_network import DeepVARTrainingNetwork, DeepVARPredictionNetwork class DeepVAREstimator(PTSEstimator): def __init__( self, + input_size: int, freq: str, prediction_length: int, target_dim: int, @@ -44,7 +45,7 @@ class DeepVAREstimator(PTSEstimator): context_length: Optional[int] = None, num_layers: int = 2, num_cells: int = 40, - cell_type: str = "lstm", + cell_type: str = "LSTM", num_parallel_samples: int = 100, dropout_rate: float = 0.1, cardinality: List[int] = [1], @@ -73,6 +74,7 @@ class DeepVAREstimator(PTSEstimator): dim=target_dim, rank=rank ) + self.input_size = input_size self.prediction_length = prediction_length self.target_dim = target_dim self.num_layers = num_layers @@ -180,6 +182,7 @@ class DeepVAREstimator(PTSEstimator): def create_training_network(self, device: torch.device) -> DeepVARTrainingNetwork: return DeepVARTrainingNetwork( + input_size=self.input_size, target_dim=self.target_dim, num_layers=self.num_layers, num_cells=self.num_cells, @@ -199,10 +202,11 @@ class DeepVAREstimator(PTSEstimator): def create_predictor( self, transformation: Transformation, - trained_network: nn.Module, + trained_network: DeepVARTrainingNetwork, device: torch.device, ) -> Predictor: prediction_network = DeepVARPredictionNetwork( + input_size=self.input_size, target_dim=self.target_dim, num_parallel_samples=self.num_parallel_samples, num_layers=self.num_layers, diff --git a/pts/model/deepvar/deepvar_network.py b/pts/model/deepvar/deepvar_network.py index f82ca36..971cd9a 100644 --- a/pts/model/deepvar/deepvar_network.py +++ b/pts/model/deepvar/deepvar_network.py @@ -6,12 +6,586 @@ from torch.distributions import Distribution import numpy as np -from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder +from pts.modules import DistributionOutput, MeanScaler, NOPScaler +from pts.model import weighted_average class DeepVARTrainingNetwork(nn.Module): - pass + 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, + dropout_rate: float, + lags_seq: List[int], + target_dim: int, + conditioning_length: int, + cardinality: List[int] = [1], + embedding_dimension: int = 1, + scaling: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + 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.target_dim = target_dim + self.scaling = scaling + self.target_dim_sample = target_dim + self.conditioning_length = conditioning_length + + assert len(set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!" + lags_seq.sort() + + self.lags_seq = lags_seq + + self.distr_output = distr_output + + self.target_dim = target_dim + + 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.proj_dist_args = distr_output.get_args_proj() + + self.embed_dim = 1 + self.embed = nn.Embedding( + num_embeddings=self.target_dim, embedding_dim=self.embed_dim + ) + + if scaling: + self.scaler = MeanScaler(keepdims=True) + else: + self.scaler = NOPScaler(keepdims=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 + the sequence from which lagged subsequences should be extracted. + Shape: (N, T, C). + sequence_length + length of sequence in the T (time) dimension (axis = 1). + indices + list of lag indices to be used. + subsequences_length + 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, :]. + """ + # we must have: history_length + begin_index >= 0 + # that is: history_length - lag_index - sequence_length >= 0 + # hence the following assert + assert max(indices) + subsequences_length <= sequence_length, ( + f"lags cannot go further than history length, found lag " + f"{max(indices)} 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, ...].unsqueeze(1)) + return torch.cat(*lagged_values, dim=1).permute(0, 2, 3, 1) + + def unroll( + self, + lags: torch.Tensor, + scale: torch.Tensor, + time_feat: torch.Tensor, + target_dimension_indicator: torch.Tensor, + unroll_length: int, + begin_state: Optional[Union[List[torch.Tensor], torch.Tensor]] = None, + ) -> Tuple[ + torch.Tensor, + Union[List[torch.Tensor], torch.Tensor], + torch.Tensor, + torch.Tensor, + ]: + + # (batch_size, sub_seq_len, target_dim, num_lags) + lags_scaled = lags / scale.unsqueeze(-1) + + # assert_shape( + # lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)), + # ) + + input_lags = lags_scaled.reshape( + (-1, unroll_length, len(self.lags_seq) * self.target_dim) + ) + + # (batch_size, target_dim, embed_dim) + index_embeddings = self.embed(target_dimension_indicator) + # assert_shape(index_embeddings, (-1, self.target_dim, self.embed_dim)) + + # (batch_size, seq_len, target_dim * embed_dim) + repeated_index_embeddings = ( + index_embeddings.unsqueeze(1) + .expand(-1, unroll_length, -1) + .reshape((-1, unroll_length, self.target_dim * self.embed_dim)) + ) + # repeated_index_embeddings = ( + # index_embeddings.expand_dims(axis=1) + # .repeat(axis=1, repeats=unroll_length) + # .reshape((-1, unroll_length, self.target_dim * self.embed_dim)) + # ) + + # (batch_size, sub_seq_len, input_dim) + inputs = torch.cat((input_lags, repeated_index_embeddings, time_feat), dim=-1) + + # unroll encoder + outputs, state = self.rnn(inputs, begin_state) + # inputs=inputs, + # length=unroll_length, + # layout="NTC", + # merge_outputs=True, + # begin_state=begin_state, + # ) + + # assert_shape(outputs, (-1, unroll_length, self.num_cells)) + # for s in state: + # assert_shape(s, (-1, self.num_cells)) + + # assert_shape( + # lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)), + # ) + + return outputs, state, lags_scaled, inputs + + def unroll_encoder( + self, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: Optional[torch.Tensor], + future_target_cdf: Optional[torch.Tensor], + target_dimension_indicator: torch.Tensor, + ) -> Tuple[ + torch.Tensor, + Union[List[torch.Tensor], torch.Tensor], + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """ + Unrolls the RNN encoder over past and, if present, future data. + Returns outputs and state of the encoder, plus the scale of + past_target_cdf and a vector of static features that was constructed + and fed as input to the encoder. All tensor arguments should have NTC + layout. + + Parameters + ---------- + past_time_feat + Past time features (batch_size, history_length, num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + future_target_cdf + Future marginal CDF transformed target values (batch_size, + prediction_length, target_dim) + target_dimension_indicator + Dimensionality of the time series (batch_size, target_dim) + + Returns + ------- + outputs + RNN outputs (batch_size, seq_len, num_cells) + states + RNN states. Nested list with (batch_size, num_cells) tensors with + dimensions target_dim x num_layers x (batch_size, num_cells) + scale + Mean scales for the time series (batch_size, 1, target_dim) + lags_scaled + Scaled lags(batch_size, sub_seq_len, target_dim, num_lags) + inputs + inputs to the RNN + + """ + + past_observed_values = torch.min( + past_observed_values, past_is_pad.unsqueeze(-1) + ) + + if future_time_feat is None or future_target_cdf is None: + time_feat = past_time_feat[:, -self.context_length :, ...] + sequence = past_target_cdf + sequence_length = self.history_length + subsequences_length = self.context_length + else: + time_feat = torch.cat( + (past_time_feat[:, -self.context_length :, ...], future_time_feat), + dim=1, + ) + sequence = torch.cat((past_target_cdf, future_target_cdf), dim=1) + sequence_length = self.history_length + self.prediction_length + subsequences_length = self.context_length + self.prediction_length + + # (batch_size, sub_seq_len, target_dim, num_lags) + lags = self.get_lagged_subsequences( + sequence=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_dim) + _, scale = self.scaler( + past_target_cdf[:, -self.context_length :, ...], + past_observed_values[:, -self.context_length : ...,], + ) + + outputs, states, lags_scaled, inputs = self.unroll( + lags=lags, + scale=scale, + time_feat=time_feat, + target_dimension_indicator=target_dimension_indicator, + unroll_length=subsequences_length, + begin_state=None, + ) + + return outputs, states, scale, lags_scaled, inputs + + def distr( + self, rnn_outputs: torch.Tensor, scale: torch.Tensor, + ): + """ + Returns the distribution of DeepVAR with respect to the RNN outputs. + + Parameters + ---------- + rnn_outputs + Outputs of the unrolled RNN (batch_size, seq_len, num_cells) + scale + Mean scale for each time series (batch_size, 1, target_dim) + + Returns + ------- + distr + Distribution instance + distr_args + Distribution arguments + """ + distr_args = self.proj_dist_args(rnn_outputs) + + # compute likelihood of target given the predicted parameters + distr = self.distr_output.distribution(distr_args, scale=scale) + + return distr, distr_args + + def forward( + self, + target_dimension_indicator: torch.Tensor, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: torch.Tensor, + future_target_cdf: torch.Tensor, + future_observed_values: torch.Tensor, + ) -> Tuple[torch.Tensor, ...]: + """ + Computes the loss for training DeepVAR, all inputs tensors representing + time series have NTC layout. + + Parameters + ---------- + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + past_time_feat + Dynamic features of past time series (batch_size, history_length, + num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + future_target_cdf + Future marginal CDF transformed target values (batch_size, + prediction_length, target_dim) + future_observed_values + Indicator whether or not the future values were observed + (batch_size, prediction_length, target_dim) + + Returns + ------- + distr + Loss with shape (batch_size, 1) + likelihoods + Likelihoods for each time step + (batch_size, context + prediction_length, 1) + distr_args + Distribution arguments (context + prediction_length, + number_of_arguments) + """ + + seq_len = self.context_length + self.prediction_length + + # unroll the decoder in "training mode", i.e. by providing future data + # as well + rnn_outputs, _, scale, lags_scaled, inputs = self.unroll_encoder( + past_time_feat=past_time_feat, + past_target_cdf=past_target_cdf, + past_observed_values=past_observed_values, + past_is_pad=past_is_pad, + future_time_feat=future_time_feat, + future_target_cdf=future_target_cdf, + target_dimension_indicator=target_dimension_indicator, + ) + + # put together target sequence + # (batch_size, seq_len, target_dim) + target = torch.cat( + (past_target_cdf[:, -self.context_length :, ...], future_target_cdf), dim=1, + ) + + # assert_shape(target, (-1, seq_len, self.target_dim)) + + distr, distr_args = self.distr(rnn_outputs=rnn_outputs, scale=scale,) + + # we sum the last axis to have the same shape for all likelihoods + # (batch_size, subseq_length, 1) + likelihoods = -distr.log_prob(target).unsqueeze(-1) + + # assert_shape(likelihoods, (-1, seq_len, 1)) + + past_observed_values = torch.min( + past_observed_values, 1 - past_is_pad.unsqueeze(-1) + ) + + # (batch_size, subseq_length, target_dim) + observed_values = torch.cat( + ( + past_observed_values[:, -self.context_length :, ...], + future_observed_values, + ), + dim=1, + ) + + # mask the loss at one time step if one or more observations is missing + # in the target dimensions (batch_size, subseq_length, 1) + loss_weights = observed_values.min(dim=-1, keepdims=True) + + # assert_shape(loss_weights, (-1, seq_len, 1)) + + loss = weighted_average(x=likelihoods, weights=loss_weights, dim=1) + + # assert_shape(loss, (-1, -1, 1)) + + self.distribution = distr + + return (loss, likelihoods) + distr_args class DeepVARPredictionNetwork(DeepVARTrainingNetwork): - pass + def __init__(self, num_parallel_samples: int, **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, + past_target_cdf: torch.Tensor, + target_dimension_indicator: torch.Tensor, + time_feat: torch.Tensor, + scale: torch.Tensor, + begin_states: Union[List[torch.Tensor], torch.Tensor], + ) -> torch.Tensor: + """ + Computes sample paths by unrolling the RNN starting with a initial + input and state. + + Parameters + ---------- + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + time_feat + Dynamic features of future time series (batch_size, history_length, + num_features) + scale + Mean scale for each time series (batch_size, 1, target_dim) + begin_states + List of initial states for the RNN layers (batch_size, num_cells) + Returns + -------- + sample_paths : Tensor + A tensor containing sampled paths. Shape: (1, num_sample_paths, + prediction_length, target_dim). + """ + + def repeat(tensor): + return tensor.repeat_interleave(repeats=self.num_parallel_samples, dim=0) + + # blows-up the dimension of each tensor to + # batch_size * self.num_sample_paths for increasing parallelism + repeated_past_target_cdf = repeat(past_target_cdf) + repeated_time_feat = repeat(time_feat) + repeated_scale = repeat(scale) + repeated_target_dimension_indicator = repeat(target_dimension_indicator) + + # slight difference for GPVAR and DeepVAR, in GPVAR, its a list + if self.cell_type == "LSTM": + repeated_states = [repeat(s) for s in begin_states] + else: + repeated_states = repeat(begin_states) + + 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): + lags = self.get_lagged_subsequences( + sequence=repeated_past_target_cdf, + sequence_length=self.history_length + k, + indices=self.shifted_lags, + subsequences_length=1, + ) + + rnn_outputs, repeated_states, lags_scaled, inputs = self.unroll( + begin_state=repeated_states, + lags=lags, + scale=repeated_scale, + time_feat=repeated_time_feat[:, k : k + 1, ...], + target_dimension_indicator=repeated_target_dimension_indicator, + unroll_length=1, + ) + + distr, distr_args = self.distr( + rnn_outputs=rnn_outputs, scale=repeated_scale, + ) + + # (batch_size, 1, target_dim) + new_samples = distr.sample() + + # (batch_size, seq_len, target_dim) + future_samples.append(new_samples) + repeated_past_target_cdf = torch.cat( + (repeated_past_target_cdf, new_samples), dim=1 + ) + + # (batch_size * num_samples, prediction_length, target_dim) + samples = torch.cat(future_samples, dim=1) + + # (batch_size, num_samples, prediction_length, target_dim) + return samples.reshape( + (-1, self.num_parallel_samples, self.prediction_length, self.target_dim,) + ) + + def forward( + self, + target_dimension_indicator: torch.Tensor, + past_time_feat: torch.Tensor, + past_target_cdf: torch.Tensor, + past_observed_values: torch.Tensor, + past_is_pad: torch.Tensor, + future_time_feat: torch.Tensor, + ) -> torch.Tensor: + """ + Predicts samples given the trained DeepVAR model. + All tensors should have NTC layout. + Parameters + ---------- + target_dimension_indicator + Indices of the target dimension (batch_size, target_dim) + past_time_feat + Dynamic features of past time series (batch_size, history_length, + num_features) + past_target_cdf + Past marginal CDF transformed target values (batch_size, + history_length, target_dim) + past_observed_values + Indicator whether or not the values were observed (batch_size, + history_length, target_dim) + past_is_pad + Indicator whether the past target values have been padded + (batch_size, history_length) + future_time_feat + Future time features (batch_size, prediction_length, num_features) + + Returns + ------- + sample_paths : Tensor + A tensor containing sampled paths (1, num_sample_paths, + prediction_length, target_dim). + + """ + + # mark padded data as unobserved + # (batch_size, target_dim, seq_len) + past_observed_values = torch.min( + past_observed_values, 1 - past_is_pad.unsqueeze(-1) + ) + + # unroll the decoder in "prediction mode", i.e. with past data only + _, state, scale, _, inputs = self.unroll_encoder( + past_time_feat=past_time_feat, + past_target_cdf=past_target_cdf, + past_observed_values=past_observed_values, + past_is_pad=past_is_pad, + future_time_feat=None, + future_target_cdf=None, + target_dimension_indicator=target_dimension_indicator, + ) + + return self.sampling_decoder( + past_target_cdf=past_target_cdf, + target_dimension_indicator=target_dimension_indicator, + time_feat=future_time_feat, + scale=scale, + begin_states=state, + )