diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py index dabf715..68460c1 100644 --- a/pts/model/deepar/deepar_network.py +++ b/pts/model/deepar/deepar_network.py @@ -281,3 +281,179 @@ class DeepARTrainingNetwork(DeepARNetwork): weighted_loss = self.weighted_average(loss, loss_weights) return weighted_loss, loss + +class DeepARPredictionNetwork(DeepARNetwork): + 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, + time_feat: torch.Tensor, + scale: torch.Tensor, + begin_states: List, + ) -> torch.Tensor: + """ + Computes sample paths by unrolling the LSTM 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 + list of initial states for the LSTM layers. + the shape of each tensor of the list should be (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( + repeats=self.num_parallel_samples, axis=0 + ) + repeated_time_feat = time_feat.repeat( + repeats=self.num_parallel_samples, axis=0 + ) + repeated_static_feat = static_feat.repeat( + repeats=self.num_parallel_samples, axis=0 + ).expand_dims(axis=1) + repeated_scale = scale.repeat( + repeats=self.num_parallel_samples, axis=0 + ) + repeated_states = [ + s.repeat(repeats=self.num_parallel_samples, axis=0) + for s in 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): + # (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, + ) + + # (batch_size * num_samples, 1, *target_shape, num_lags) + lags_scaled = F.broadcast_div( + lags, repeated_scale.expand_dims(axis=-1) + ) + + # from (batch_size * num_samples, 1, *target_shape, num_lags) + # to (batch_size * num_samples, 1, prod(target_shape) * num_lags) + input_lags = F.reshape( + data=lags_scaled, + shape=(-1, 1, prod(self.target_shape) * len(self.lags_seq)), + ) + + # (batch_size * num_samples, 1, prod(target_shape) * num_lags + num_time_features + num_static_features) + decoder_input = F.concat( + input_lags, + repeated_time_feat.slice_axis(axis=1, begin=k, end=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.unroll( + inputs=decoder_input, + length=1, + begin_state=repeated_states, + layout="NTC", + merge_outputs=True, + ) + + distr_args = self.proj_distr_args(rnn_outputs) + + # 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(dtype=self.dtype) + + # (batch_size * num_samples, seq_len, *target_shape) + repeated_past_target = F.concat( + repeated_past_target, new_samples, dim=1 + ) + future_samples.append(new_samples) + + # (batch_size * num_samples, prediction_length, *target_shape) + samples = F.concat(*future_samples, dim=1) + + # (batch_size, num_samples, prediction_length, *target_shape) + return samples.reshape( + shape=( + (-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) + ) -> torch.Tensor: + """ + Predicts samples, all tensors should have NTC layout. + Parameters + ---------- + F + 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, 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, + future_time_feat=None, + future_target=None, + ) + + return self.sampling_decoder( + past_target=past_target, + time_feat=future_time_feat, + static_feat=static_feat, + scale=scale, + begin_states=state, + )