mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-07-08 03:46:12 +08:00
Nick
ccc35afb31
* GluonTS import updates * drop freq argument see https://github.com/awslabs/gluon-ts/pull/1997
458 lines
17 KiB
Python
458 lines
17 KiB
Python
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.distributions.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 DeepARNetwork(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,
|
|
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
|
|
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)
|
|
|
|
self.embedder = FeatureEmbedder(
|
|
cardinalities=cardinality, embedding_dims=embedding_dimension
|
|
)
|
|
|
|
if scaling:
|
|
self.scaler = MeanScaler(keepdim=True)
|
|
else:
|
|
self.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)
|
|
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:
|
|
time_feat = past_time_feat[
|
|
:, self.history_length - self.context_length :, ...
|
|
]
|
|
sequence = past_target
|
|
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)
|
|
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,
|
|
)
|
|
|
|
# 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 :, ...],
|
|
)
|
|
|
|
# (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(),
|
|
),
|
|
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)
|
|
|
|
# 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))
|
|
)
|
|
|
|
# (batch_size, sub_seq_len, input_dim)
|
|
inputs = torch.cat((input_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, static_feat
|
|
|
|
|
|
class DeepARTrainingNetwork(DeepARNetwork):
|
|
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 = weighted_average(loss, weights=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: 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_static_feat = static_feat.repeat_interleave(
|
|
repeats=self.num_parallel_samples, dim=0
|
|
).unsqueeze(1)
|
|
repeated_scale = 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,
|
|
)
|
|
|
|
# (batch_size * num_samples, 1, *target_shape, num_lags)
|
|
lags_scaled = lags / repeated_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))
|
|
)
|
|
|
|
# (batch_size * num_samples, 1, prod(target_shape) * num_lags + num_time_features + num_static_features)
|
|
decoder_input = torch.cat(
|
|
(input_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)
|
|
|
|
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()
|
|
|
|
# (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)
|
|
) -> 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, 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,
|
|
)
|