diff --git a/pts/model/deepar/deepar_network.py b/pts/model/deepar/deepar_network.py index 4f558f7..c9e59df 100644 --- a/pts/model/deepar/deepar_network.py +++ b/pts/model/deepar/deepar_network.py @@ -1,7 +1,8 @@ -from typing import List +from typing import List, Optional, Tuple import torch import torch.nn as nn +from torch.distributions import Distribution import numpy as np @@ -189,4 +190,76 @@ class DeepARNetwork(nn.Module): return outputs, state, scale, static_feat class DeepARTrainingNetwork(DeepARNetwork): - pass \ No newline at end of file + 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 + ) -> Distribution: + rnn_outputs, _, 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, + ) + + distr_args = self.proj_distr_args(rnn_outputs) + + return 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 + ) -> torch.Tensor: + 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, + ) + + # 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) + + # (batch_size, seq_len) + loss = -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 = self.weighted_average(loss, loss_weights) + + return weighted_loss, loss