mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-08-11 11:24:31 +08:00
Temporal fusion transformer (#40)
* initial tft model * fixes * fix FeatureProjector * fixed attention module * added quantile loss output * added predict * added tft_transform from gluonts * requires tensor_split * comment out key_padding_mask give nans on GPU * added example
This commit is contained in:
Kashif Rasul
File diff suppressed because one or more lines are too long
@@ -0,0 +1,5 @@
|
||||
from .tft_estimator import TemporalFusionTransformerEstimator
|
||||
from .tft_network import (
|
||||
TemporalFusionTransformerTrainingNetwork,
|
||||
TemporalFusionTransformerPredictionNetwork,
|
||||
)
|
||||
@@ -0,0 +1,385 @@
|
||||
from typing import List, Optional, Dict
|
||||
from itertools import chain
|
||||
|
||||
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.torch.modules.distribution_output import DistributionOutput
|
||||
from gluonts.torch.support.util import copy_parameters
|
||||
from gluonts.torch.model.predictor import PyTorchPredictor
|
||||
from gluonts.model.predictor import Predictor
|
||||
from gluonts.model.forecast_generator import QuantileForecastGenerator
|
||||
from gluonts.transform import (
|
||||
Transformation,
|
||||
Chain,
|
||||
ValidationSplitSampler,
|
||||
TestSplitSampler,
|
||||
ExpectedNumInstanceSampler,
|
||||
AddAgeFeature,
|
||||
AsNumpyArray,
|
||||
AddObservedValuesIndicator,
|
||||
AddTimeFeatures,
|
||||
VstackFeatures,
|
||||
SetField,
|
||||
)
|
||||
|
||||
from pts import Trainer
|
||||
from pts.model.utils import get_module_forward_input_names
|
||||
from pts.model import PyTorchEstimator
|
||||
|
||||
from .tft_network import (
|
||||
TemporalFusionTransformerPredictionNetwork,
|
||||
TemporalFusionTransformerTrainingNetwork,
|
||||
)
|
||||
from .tft_transform import BroadcastTo, TFTInstanceSplitter
|
||||
|
||||
|
||||
def _default_feat_args(dims_or_cardinalities: List[int]):
|
||||
if dims_or_cardinalities:
|
||||
return dims_or_cardinalities
|
||||
return [1]
|
||||
|
||||
|
||||
class TemporalFusionTransformerEstimator(PyTorchEstimator):
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
freq: str,
|
||||
prediction_length: int,
|
||||
context_length: Optional[int] = None,
|
||||
dropout_rate: float = 0.1,
|
||||
embed_dim: int = 32,
|
||||
num_heads: int = 4,
|
||||
num_outputs: int = 3,
|
||||
variable_dim: Optional[int] = None,
|
||||
time_features: List[TimeFeature] = [],
|
||||
static_cardinalities: Dict[str, int] = {},
|
||||
dynamic_cardinalities: Dict[str, int] = {},
|
||||
static_feature_dims: Dict[str, int] = {},
|
||||
dynamic_feature_dims: Dict[str, int] = {},
|
||||
past_dynamic_features: List[str] = [],
|
||||
trainer: Trainer = Trainer(),
|
||||
) -> None:
|
||||
super().__init__(trainer=trainer)
|
||||
|
||||
self.freq = freq
|
||||
self.prediction_length = prediction_length
|
||||
self.context_length = context_length or prediction_length
|
||||
|
||||
# MultiheadAttention
|
||||
self.dropout_rate = dropout_rate
|
||||
self.embed_dim = embed_dim
|
||||
self.num_heads = num_heads
|
||||
|
||||
self.num_outputs = num_outputs
|
||||
self.variable_dim = variable_dim or embed_dim
|
||||
|
||||
if not time_features:
|
||||
self.time_features = time_features_from_frequency_str(self.freq)
|
||||
else:
|
||||
self.time_features = time_features
|
||||
self.static_cardinalities = static_cardinalities
|
||||
self.dynamic_cardinalities = dynamic_cardinalities
|
||||
self.static_feature_dims = static_feature_dims
|
||||
self.dynamic_feature_dims = dynamic_feature_dims
|
||||
self.past_dynamic_features = past_dynamic_features
|
||||
|
||||
self.past_dynamic_cardinalities = {}
|
||||
self.past_dynamic_feature_dims = {}
|
||||
for name in self.past_dynamic_features:
|
||||
if name in self.dynamic_cardinalities:
|
||||
self.past_dynamic_cardinalities[name] = self.dynamic_cardinalities.pop(
|
||||
name
|
||||
)
|
||||
elif name in self.dynamic_feature_dims:
|
||||
self.past_dynamic_feature_dims[name] = self.dynamic_feature_dims.pop(
|
||||
name
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Feature name {name} is not provided in feature dicts"
|
||||
)
|
||||
|
||||
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:
|
||||
transforms = (
|
||||
[AsNumpyArray(field=FieldName.TARGET, expected_ndim=1)]
|
||||
+ (
|
||||
[
|
||||
AsNumpyArray(field=name, expected_ndim=1)
|
||||
for name in self.static_cardinalities.keys()
|
||||
]
|
||||
)
|
||||
+ [
|
||||
AsNumpyArray(field=name, expected_ndim=1)
|
||||
for name in chain(
|
||||
self.static_feature_dims.keys(),
|
||||
self.dynamic_cardinalities.keys(),
|
||||
)
|
||||
]
|
||||
+ [
|
||||
AsNumpyArray(field=name, expected_ndim=2)
|
||||
for name in self.dynamic_feature_dims.keys()
|
||||
]
|
||||
+ [
|
||||
AddObservedValuesIndicator(
|
||||
target_field=FieldName.TARGET,
|
||||
output_field=FieldName.OBSERVED_VALUES,
|
||||
),
|
||||
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,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
if self.static_cardinalities:
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
output_field=FieldName.FEAT_STATIC_CAT,
|
||||
input_fields=list(self.static_cardinalities.keys()),
|
||||
h_stack=True,
|
||||
)
|
||||
)
|
||||
else:
|
||||
transforms.extend(
|
||||
[
|
||||
SetField(
|
||||
output_field=FieldName.FEAT_STATIC_CAT,
|
||||
value=[0],
|
||||
),
|
||||
AsNumpyArray(
|
||||
field=FieldName.FEAT_STATIC_CAT, expected_ndim=1, dtype=np.long
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
if self.static_feature_dims:
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
output_field=FieldName.FEAT_STATIC_REAL,
|
||||
input_fields=list(self.static_feature_dims.keys()),
|
||||
h_stack=True,
|
||||
)
|
||||
)
|
||||
else:
|
||||
transforms.extend(
|
||||
[
|
||||
SetField(
|
||||
output_field=FieldName.FEAT_STATIC_REAL,
|
||||
value=[0.0],
|
||||
),
|
||||
AsNumpyArray(field=FieldName.FEAT_STATIC_REAL, expected_ndim=1),
|
||||
]
|
||||
)
|
||||
|
||||
if self.dynamic_cardinalities:
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
output_field=FieldName.FEAT_DYNAMIC_CAT,
|
||||
input_fields=list(self.dynamic_cardinalities.keys()),
|
||||
)
|
||||
)
|
||||
else:
|
||||
transforms.extend(
|
||||
[
|
||||
SetField(
|
||||
output_field=FieldName.FEAT_DYNAMIC_CAT,
|
||||
value=[[0]],
|
||||
),
|
||||
AsNumpyArray(
|
||||
field=FieldName.FEAT_DYNAMIC_CAT,
|
||||
expected_ndim=2,
|
||||
dtype=np.long,
|
||||
),
|
||||
BroadcastTo(
|
||||
field=FieldName.FEAT_DYNAMIC_CAT,
|
||||
ext_length=self.prediction_length,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
input_fields = [FieldName.FEAT_TIME, FieldName.FEAT_AGE]
|
||||
if self.dynamic_feature_dims:
|
||||
input_fields += list(self.dynamic_feature_dims.keys())
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
input_fields=input_fields,
|
||||
output_field=FieldName.FEAT_DYNAMIC_REAL,
|
||||
)
|
||||
)
|
||||
|
||||
if self.past_dynamic_cardinalities:
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
output_field=FieldName.PAST_FEAT_DYNAMIC + "_cat",
|
||||
input_fields=list(self.past_dynamic_cardinalities.keys()),
|
||||
)
|
||||
)
|
||||
else:
|
||||
transforms.extend(
|
||||
[
|
||||
SetField(
|
||||
output_field=FieldName.PAST_FEAT_DYNAMIC + "_cat",
|
||||
value=[[0]],
|
||||
),
|
||||
AsNumpyArray(
|
||||
field=FieldName.PAST_FEAT_DYNAMIC + "_cat",
|
||||
expected_ndim=2,
|
||||
dtype=np.long,
|
||||
),
|
||||
BroadcastTo(field=FieldName.PAST_FEAT_DYNAMIC + "_cat"),
|
||||
]
|
||||
)
|
||||
|
||||
if self.past_dynamic_feature_dims:
|
||||
transforms.append(
|
||||
VstackFeatures(
|
||||
output_field=FieldName.PAST_FEAT_DYNAMIC_REAL,
|
||||
input_fields=list(self.past_dynamic_feature_dims.keys()),
|
||||
)
|
||||
)
|
||||
else:
|
||||
transforms.extend(
|
||||
[
|
||||
SetField(
|
||||
output_field=FieldName.PAST_FEAT_DYNAMIC_REAL,
|
||||
value=[[0.0]],
|
||||
),
|
||||
AsNumpyArray(
|
||||
field=FieldName.PAST_FEAT_DYNAMIC_REAL, expected_ndim=2
|
||||
),
|
||||
BroadcastTo(field=FieldName.PAST_FEAT_DYNAMIC_REAL),
|
||||
]
|
||||
)
|
||||
|
||||
return Chain(transforms)
|
||||
|
||||
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]
|
||||
|
||||
ts_fields = [FieldName.FEAT_DYNAMIC_CAT, FieldName.FEAT_DYNAMIC_REAL]
|
||||
past_ts_fields = [
|
||||
FieldName.PAST_FEAT_DYNAMIC + "_cat",
|
||||
FieldName.PAST_FEAT_DYNAMIC_REAL,
|
||||
]
|
||||
|
||||
return TFTInstanceSplitter(
|
||||
instance_sampler=instance_sampler,
|
||||
past_length=self.context_length,
|
||||
future_length=self.prediction_length,
|
||||
time_series_fields=ts_fields,
|
||||
past_time_series_fields=past_ts_fields,
|
||||
)
|
||||
|
||||
def create_training_network(
|
||||
self, device: torch.device
|
||||
) -> TemporalFusionTransformerTrainingNetwork:
|
||||
network = TemporalFusionTransformerTrainingNetwork(
|
||||
context_length=self.context_length,
|
||||
prediction_length=self.prediction_length,
|
||||
variable_dim=self.variable_dim,
|
||||
embed_dim=self.embed_dim,
|
||||
num_heads=self.num_heads,
|
||||
num_outputs=self.num_outputs,
|
||||
dropout=self.dropout_rate,
|
||||
d_past_feat_dynamic_real=_default_feat_args(
|
||||
list(self.past_dynamic_feature_dims.values())
|
||||
),
|
||||
c_past_feat_dynamic_cat=_default_feat_args(
|
||||
list(self.past_dynamic_cardinalities.values())
|
||||
),
|
||||
d_feat_dynamic_real=_default_feat_args(
|
||||
[1] * len(self.time_features) + list(self.dynamic_feature_dims.values())
|
||||
),
|
||||
c_feat_dynamic_cat=_default_feat_args(
|
||||
list(self.dynamic_cardinalities.values())
|
||||
),
|
||||
d_feat_static_real=_default_feat_args(
|
||||
list(self.static_feature_dims.values()),
|
||||
),
|
||||
c_feat_static_cat=_default_feat_args(
|
||||
list(self.static_cardinalities.values()),
|
||||
),
|
||||
)
|
||||
return network.to(device)
|
||||
|
||||
def create_predictor(
|
||||
self,
|
||||
transformation: Transformation,
|
||||
trained_network: TemporalFusionTransformerTrainingNetwork,
|
||||
device: torch.device,
|
||||
) -> Predictor:
|
||||
|
||||
prediction_network = TemporalFusionTransformerPredictionNetwork(
|
||||
context_length=self.context_length,
|
||||
prediction_length=self.prediction_length,
|
||||
variable_dim=self.variable_dim,
|
||||
embed_dim=self.embed_dim,
|
||||
num_heads=self.num_heads,
|
||||
num_outputs=self.num_outputs,
|
||||
dropout=self.dropout_rate,
|
||||
d_past_feat_dynamic_real=_default_feat_args(
|
||||
list(self.past_dynamic_feature_dims.values())
|
||||
),
|
||||
c_past_feat_dynamic_cat=_default_feat_args(
|
||||
list(self.past_dynamic_cardinalities.values())
|
||||
),
|
||||
d_feat_dynamic_real=_default_feat_args(
|
||||
[1] * len(self.time_features) + list(self.dynamic_feature_dims.values())
|
||||
),
|
||||
c_feat_dynamic_cat=_default_feat_args(
|
||||
list(self.dynamic_cardinalities.values())
|
||||
),
|
||||
d_feat_static_real=_default_feat_args(
|
||||
list(self.static_feature_dims.values()),
|
||||
),
|
||||
c_feat_static_cat=_default_feat_args(
|
||||
list(self.static_cardinalities.values()),
|
||||
),
|
||||
).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,
|
||||
forecast_generator=QuantileForecastGenerator(
|
||||
quantiles=[str(q) for q in prediction_network.quantiles],
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,295 @@
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from pts.modules import FeatureEmbedder as BaseFeatureEmbedder
|
||||
|
||||
|
||||
class FeatureProjector(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
feature_dims: List[int],
|
||||
embedding_dims: List[int],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.__num_features = len(feature_dims)
|
||||
if self.__num_features > 1:
|
||||
self.feature_dims = (
|
||||
feature_dims[0:1] + np.cumsum(feature_dims)[:-1].tolist()
|
||||
)
|
||||
else:
|
||||
self.feature_dims = feature_dims
|
||||
|
||||
self._projector = nn.ModuleList(
|
||||
[
|
||||
nn.Linear(in_features=in_feature, out_features=out_features)
|
||||
for in_feature, out_features in zip(self.feature_dims, embedding_dims)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, features: torch.Tensor) -> List[torch.Tensor]:
|
||||
if self.__num_features > 1:
|
||||
real_feature_slices = torch.tensor_split(
|
||||
features, self.feature_dims[1:], dim=-1
|
||||
)
|
||||
else:
|
||||
real_feature_slices = [features]
|
||||
|
||||
return [
|
||||
proj(real_feature_slice)
|
||||
for proj, real_feature_slice in zip(self._projector, real_feature_slices)
|
||||
]
|
||||
|
||||
|
||||
class FeatureEmbedder(BaseFeatureEmbedder):
|
||||
def forward(self, features: torch.Tensor) -> List[torch.Tensor]:
|
||||
concat_features = super(FeatureEmbedder, self).forward(features=features)
|
||||
|
||||
if self.__num_features > 1:
|
||||
features = torch.chunk(concat_features, self.__num_features, dim=-1)
|
||||
else:
|
||||
features = [concat_features]
|
||||
|
||||
return features
|
||||
|
||||
|
||||
class GatedLinearUnit(nn.Module):
|
||||
def __init__(self, dim: int = -1, nonlinear: bool = True):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.nonlinear = nonlinear
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
val, gate = torch.chunk(x, 2, dim=self.dim)
|
||||
if self.nonlinear:
|
||||
val = torch.tanh(val)
|
||||
return torch.sigmoid(gate) * val
|
||||
|
||||
|
||||
class GatedResidualNetwork(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_hidden: int,
|
||||
d_input: Optional[int] = None,
|
||||
d_output: Optional[int] = None,
|
||||
d_static: Optional[int] = None,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
d_input = d_input or d_hidden
|
||||
d_static = d_static or 0
|
||||
if d_output is None:
|
||||
d_output = d_input
|
||||
self.add_skip = False
|
||||
else:
|
||||
if d_output != d_input:
|
||||
self.add_skip = True
|
||||
self.skip_proj = nn.Linear(in_features=d_input, out_features=d_output)
|
||||
else:
|
||||
self.add_skip = False
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(in_features=d_input + d_static, out_features=d_hidden),
|
||||
nn.ELU(),
|
||||
nn.Linear(in_features=d_hidden, out_features=d_hidden),
|
||||
nn.Dropout(p=dropout),
|
||||
nn.Linear(in_features=d_hidden, out_features=d_output * 2),
|
||||
GatedLinearUnit(nonlinear=False),
|
||||
)
|
||||
|
||||
self.lnorm = nn.LayerNorm(d_output)
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, c: Optional[torch.Tensor] = None
|
||||
) -> torch.Tensor:
|
||||
if self.add_skip:
|
||||
skip = self.skip_proj(x)
|
||||
else:
|
||||
skip = x
|
||||
|
||||
if c is not None:
|
||||
x = torch.cat((x, c), dim=-1)
|
||||
x = self.mlp(x)
|
||||
x = self.lnorm(x + skip)
|
||||
return x
|
||||
|
||||
|
||||
class VariableSelectionNetwork(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_hidden: int,
|
||||
n_vars: int,
|
||||
dropout: float = 0.0,
|
||||
add_static: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.weight_network = GatedResidualNetwork(
|
||||
d_hidden=d_hidden,
|
||||
d_input=d_hidden * n_vars,
|
||||
d_output=n_vars,
|
||||
d_static=d_hidden if add_static else None,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.variable_network = nn.ModuleList(
|
||||
[
|
||||
GatedResidualNetwork(d_hidden=d_hidden, dropout=dropout)
|
||||
for _ in range(n_vars)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, variables: List[torch.Tensor], static: Optional[torch.Tensor] = None
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
flatten = torch.cat(variables, dim=-1)
|
||||
if static is not None:
|
||||
static = static.expand_as(variables[0])
|
||||
weight = self.weight_network(flatten, static)
|
||||
weight = torch.softmax(weight.unsqueeze(-2), dim=-1)
|
||||
|
||||
var_encodings = [net(var) for var, net in zip(variables, self.variable_network)]
|
||||
var_encodings = torch.stack(var_encodings, dim=-1)
|
||||
|
||||
var_encodings = torch.sum(var_encodings * weight, dim=-1)
|
||||
|
||||
return var_encodings, weight
|
||||
|
||||
|
||||
class TemporalFusionEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_input: int,
|
||||
d_hidden: int,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.encoder_lstm = nn.LSTM(
|
||||
input_size=d_input, hidden_size=d_hidden, batch_first=True
|
||||
)
|
||||
self.decoder_lstm = nn.LSTM(
|
||||
input_size=d_input, hidden_size=d_hidden, batch_first=True
|
||||
)
|
||||
|
||||
self.gate = nn.Sequential(
|
||||
nn.Linear(in_features=d_hidden, out_features=d_hidden * 2),
|
||||
GatedLinearUnit(nonlinear=False),
|
||||
)
|
||||
if d_input != d_hidden:
|
||||
self.skip_proj = nn.Linear(in_features=d_input, out_features=d_hidden)
|
||||
self.add_skip = True
|
||||
else:
|
||||
self.add_skip = False
|
||||
|
||||
self.lnorm = nn.LayerNorm(d_hidden)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
ctx_input: torch.Tensor,
|
||||
tgt_input: torch.Tensor,
|
||||
states: List[torch.Tensor],
|
||||
):
|
||||
ctx_encodings, states = self.encoder_lstm(ctx_input, states)
|
||||
|
||||
tgt_encodings, _ = self.decoder_lstm(tgt_input, states)
|
||||
|
||||
encodings = torch.cat((ctx_encodings, tgt_encodings), dim=1)
|
||||
skip = torch.cat((ctx_input, tgt_input), dim=1)
|
||||
if self.add_skip:
|
||||
skip = self.skip_proj(skip)
|
||||
encodings = self.gate(encodings)
|
||||
encodings = self.lnorm(skip + encodings)
|
||||
return encodings
|
||||
|
||||
|
||||
class TemporalFusionDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
context_length: int,
|
||||
prediction_length: int,
|
||||
d_hidden: int,
|
||||
d_var: int,
|
||||
n_head: int,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.context_length = context_length
|
||||
self.prediction_length = prediction_length
|
||||
|
||||
self.enrich = GatedResidualNetwork(
|
||||
d_hidden=d_hidden,
|
||||
d_static=d_var,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.attention = nn.MultiheadAttention(
|
||||
embed_dim=d_hidden, num_heads=n_head, dropout=dropout
|
||||
)
|
||||
|
||||
self.att_net = nn.Sequential(
|
||||
nn.Linear(in_features=d_hidden, out_features=d_hidden * 2),
|
||||
GatedLinearUnit(nonlinear=False),
|
||||
)
|
||||
self.att_lnorm = nn.LayerNorm(d_hidden)
|
||||
|
||||
self.ff_net = nn.Sequential(
|
||||
GatedResidualNetwork(d_hidden=d_hidden, dropout=dropout),
|
||||
nn.Linear(in_features=d_hidden, out_features=d_hidden * 2),
|
||||
GatedLinearUnit(nonlinear=False),
|
||||
)
|
||||
self.ff_lnorm = nn.LayerNorm(d_hidden)
|
||||
|
||||
self.register_buffer(
|
||||
"attn_mask",
|
||||
self._generate_subsequent_mask(
|
||||
prediction_length, prediction_length + context_length
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _generate_subsequent_mask(
|
||||
target_length: int, source_length: int
|
||||
) -> torch.Tensor:
|
||||
mask = (torch.triu(torch.ones(source_length, target_length)) == 1).transpose(
|
||||
0, 1
|
||||
)
|
||||
mask = (
|
||||
mask.float()
|
||||
.masked_fill(mask == 0, float("-inf"))
|
||||
.masked_fill(mask == 1, float(0.0))
|
||||
)
|
||||
return mask
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, static: torch.Tensor, mask: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
static = static.repeat((1, self.context_length + self.prediction_length, 1))
|
||||
|
||||
skip = x[:, self.context_length :, ...]
|
||||
x = self.enrich(x, static)
|
||||
|
||||
mask_pad = torch.ones_like(mask)[:, 0:1, ...]
|
||||
mask_pad = mask_pad.repeat((1, self.prediction_length))
|
||||
key_padding_mask = torch.cat((mask, mask_pad), dim=1).bool()
|
||||
|
||||
query_key_value = x.permute(1, 0, 2)
|
||||
|
||||
attn_output, _ = self.attention(
|
||||
query=query_key_value[-self.prediction_length :, ...],
|
||||
key=query_key_value,
|
||||
value=query_key_value,
|
||||
# key_padding_mask=key_padding_mask, # does not work on GPU :-(
|
||||
attn_mask=self.attn_mask,
|
||||
)
|
||||
att = self.att_net(attn_output.permute(1, 0, 2))
|
||||
|
||||
x = x[:, self.context_length :, ...]
|
||||
x = self.att_lnorm(x + att)
|
||||
x = self.ff_net(x)
|
||||
x = self.ff_lnorm(x + skip)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,346 @@
|
||||
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 pts.model import weighted_average
|
||||
|
||||
from .tft_modules import (
|
||||
FeatureProjector,
|
||||
FeatureEmbedder,
|
||||
VariableSelectionNetwork,
|
||||
GatedResidualNetwork,
|
||||
TemporalFusionEncoder,
|
||||
TemporalFusionDecoder,
|
||||
)
|
||||
from .tft_output import QuantileOutput
|
||||
|
||||
|
||||
class TemporalFusionTransformerNetwork(nn.Module):
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
context_length: int,
|
||||
prediction_length: int,
|
||||
variable_dim: int,
|
||||
embed_dim: int,
|
||||
num_heads: int,
|
||||
num_outputs: int,
|
||||
d_past_feat_dynamic_real: List[int],
|
||||
c_past_feat_dynamic_cat: List[int],
|
||||
d_feat_dynamic_real: List[int],
|
||||
c_feat_dynamic_cat: List[int],
|
||||
d_feat_static_real: List[int],
|
||||
c_feat_static_cat: List[int],
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.context_length = context_length
|
||||
self.prediction_length = prediction_length
|
||||
self.normalize_eps = 1e-5
|
||||
|
||||
self.target_proj = nn.Linear(in_features=1, out_features=variable_dim)
|
||||
|
||||
if d_past_feat_dynamic_real:
|
||||
self.past_feat_dynamic_proj = FeatureProjector(
|
||||
feature_dims=d_past_feat_dynamic_real,
|
||||
embedding_dims=[variable_dim] * len(d_past_feat_dynamic_real),
|
||||
)
|
||||
else:
|
||||
self.past_feat_dynamic_proj = None
|
||||
|
||||
if c_past_feat_dynamic_cat:
|
||||
self.past_feat_dynamic_embed = FeatureEmbedder(
|
||||
cardinalities=c_past_feat_dynamic_cat,
|
||||
embedding_dims=[variable_dim] * len(c_past_feat_dynamic_cat),
|
||||
)
|
||||
else:
|
||||
self.past_feat_dynamic_embed = None
|
||||
|
||||
if d_feat_dynamic_real:
|
||||
self.feat_dynamic_proj = FeatureProjector(
|
||||
feature_dims=d_feat_dynamic_real,
|
||||
embedding_dims=[variable_dim] * len(d_feat_dynamic_real),
|
||||
)
|
||||
else:
|
||||
self.feat_dynamic_proj = None
|
||||
|
||||
if c_feat_dynamic_cat:
|
||||
self.feat_dynamic_embed = FeatureEmbedder(
|
||||
cardinalities=c_feat_dynamic_cat,
|
||||
embedding_dims=[variable_dim] * len(c_feat_dynamic_cat),
|
||||
)
|
||||
else:
|
||||
self.feat_dynamic_embed = None
|
||||
|
||||
if d_feat_static_real:
|
||||
self.feat_static_proj = FeatureProjector(
|
||||
feature_dims=d_feat_static_real,
|
||||
embedding_dims=[variable_dim] * len(d_feat_static_real),
|
||||
)
|
||||
else:
|
||||
self.feat_static_proj = None
|
||||
|
||||
if c_feat_static_cat:
|
||||
self.feat_static_embed = FeatureEmbedder(
|
||||
cardinalities=c_feat_static_cat,
|
||||
embedding_dims=[variable_dim] * len(c_feat_static_cat),
|
||||
)
|
||||
else:
|
||||
self.feat_static_embed = None
|
||||
|
||||
n_feat_static = len(d_feat_static_real) + len(c_feat_static_cat)
|
||||
self.static_selector = VariableSelectionNetwork(
|
||||
d_hidden=variable_dim,
|
||||
n_vars=n_feat_static,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
n_past_feat_dynamic = len(d_past_feat_dynamic_real) + len(
|
||||
c_past_feat_dynamic_cat
|
||||
)
|
||||
n_feat_dynamic = len(d_feat_dynamic_real) + len(c_feat_dynamic_cat)
|
||||
self.ctx_selector = VariableSelectionNetwork(
|
||||
d_hidden=variable_dim,
|
||||
n_vars=n_past_feat_dynamic + n_feat_dynamic + 1,
|
||||
add_static=True,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.tgt_selector = VariableSelectionNetwork(
|
||||
d_hidden=variable_dim,
|
||||
n_vars=n_feat_dynamic,
|
||||
add_static=True,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.selection = GatedResidualNetwork(
|
||||
d_hidden=variable_dim,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.enrichment = GatedResidualNetwork(
|
||||
d_hidden=variable_dim,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.state_h = GatedResidualNetwork(
|
||||
d_hidden=variable_dim,
|
||||
d_output=embed_dim,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.state_c = GatedResidualNetwork(
|
||||
d_hidden=variable_dim,
|
||||
d_output=embed_dim,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.temporal_encoder = TemporalFusionEncoder(
|
||||
d_input=variable_dim,
|
||||
d_hidden=embed_dim,
|
||||
)
|
||||
self.temporal_decoder = TemporalFusionDecoder(
|
||||
context_length=self.context_length,
|
||||
prediction_length=self.prediction_length,
|
||||
d_hidden=embed_dim,
|
||||
d_var=variable_dim,
|
||||
n_head=num_heads,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
self.quantiles = sum(
|
||||
[[i / 10, 1.0 - i / 10] for i in range(1, (num_outputs + 1) // 2)],
|
||||
[0.5],
|
||||
)
|
||||
self.output = QuantileOutput(input_size=embed_dim, quantiles=self.quantiles)
|
||||
self.output_proj = self.output.get_quantile_proj()
|
||||
self.loss = self.output.get_loss()
|
||||
|
||||
def _preprocess(
|
||||
self,
|
||||
past_target: torch.Tensor,
|
||||
past_observed_values: torch.Tensor,
|
||||
past_feat_dynamic_real: torch.Tensor,
|
||||
past_feat_dynamic_cat: torch.Tensor,
|
||||
feat_dynamic_real: torch.Tensor,
|
||||
feat_dynamic_cat: torch.Tensor,
|
||||
feat_static_real: torch.Tensor,
|
||||
feat_static_cat: torch.Tensor,
|
||||
):
|
||||
obs = past_target * past_observed_values
|
||||
count = past_observed_values.sum(dim=1, keepdim=True)
|
||||
offset = obs.sum(1, keepdim=True) / (count + self.normalize_eps)
|
||||
scale = torch.sum(obs ** 2, 1, keepdim=True) / (count + self.normalize_eps)
|
||||
scale = torch.sqrt(scale - offset ** 2)
|
||||
|
||||
past_target = (past_target - offset) / (scale + self.normalize_eps)
|
||||
past_target = past_target.unsqueeze(-1)
|
||||
|
||||
proj = self.target_proj(past_target)
|
||||
|
||||
past_covariates = []
|
||||
future_covariates = []
|
||||
static_covariates: List[torch.Tensor] = []
|
||||
|
||||
past_covariates.append(proj)
|
||||
if self.past_feat_dynamic_proj is not None:
|
||||
projs = self.past_feat_dynamic_proj(past_feat_dynamic_real)
|
||||
past_covariates.extend(projs)
|
||||
if self.past_feat_dynamic_embed is not None:
|
||||
embs = self.past_feat_dynamic_embed(past_feat_dynamic_cat)
|
||||
past_covariates.extend(embs)
|
||||
if self.feat_dynamic_proj is not None:
|
||||
projs = self.feat_dynamic_proj(feat_dynamic_real)
|
||||
for proj in projs:
|
||||
ctx_proj = proj[:, 0 : self.context_length, ...]
|
||||
tgt_proj = proj[:, self.context_length :, ...]
|
||||
past_covariates.append(ctx_proj)
|
||||
future_covariates.append(tgt_proj)
|
||||
if self.feat_dynamic_embed is not None:
|
||||
embs = self.feat_dynamic_embed(feat_dynamic_cat)
|
||||
for emb in embs:
|
||||
ctx_emb = emb[:, 0 : self.context_length, ...]
|
||||
tgt_emb = emb[:, self.context_length :, ...]
|
||||
past_covariates.append(ctx_emb)
|
||||
future_covariates.append(tgt_emb)
|
||||
|
||||
if self.feat_static_proj is not None:
|
||||
projs = self.feat_static_proj(feat_static_real)
|
||||
static_covariates.extend(projs)
|
||||
if self.feat_static_embed is not None:
|
||||
embs = self.feat_static_embed(feat_static_cat)
|
||||
static_covariates.extend(embs)
|
||||
|
||||
return (
|
||||
past_covariates,
|
||||
future_covariates,
|
||||
static_covariates,
|
||||
offset,
|
||||
scale,
|
||||
)
|
||||
|
||||
def _postprocess(
|
||||
self,
|
||||
preds: torch.Tensor,
|
||||
offset: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
offset = offset.unsqueeze(-1)
|
||||
scale = scale.unsqueeze(-1)
|
||||
preds = (preds * (scale + self.normalize_eps)) + offset
|
||||
return preds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
past_observed_values: torch.Tensor,
|
||||
past_covariates: torch.Tensor,
|
||||
future_covariates: torch.Tensor,
|
||||
static_covariates: torch.Tensor,
|
||||
):
|
||||
static_var, _ = self.static_selector(static_covariates)
|
||||
c_selection = self.selection(static_var).unsqueeze(1)
|
||||
c_enrichment = self.enrichment(static_var).unsqueeze(1)
|
||||
c_h = self.state_h(static_var)
|
||||
c_c = self.state_c(static_var)
|
||||
|
||||
ctx_input, _ = self.ctx_selector(past_covariates, c_selection)
|
||||
tgt_input, _ = self.tgt_selector(future_covariates, c_selection)
|
||||
|
||||
encoding = self.temporal_encoder(
|
||||
ctx_input, tgt_input, [c_h.unsqueeze(0), c_c.unsqueeze(0)]
|
||||
)
|
||||
decoding = self.temporal_decoder(encoding, c_enrichment, past_observed_values)
|
||||
|
||||
preds = self.output_proj(decoding)
|
||||
|
||||
return preds
|
||||
|
||||
|
||||
class TemporalFusionTransformerTrainingNetwork(TemporalFusionTransformerNetwork):
|
||||
def forward(
|
||||
self,
|
||||
past_target: torch.Tensor,
|
||||
past_observed_values: torch.Tensor,
|
||||
future_target: torch.Tensor,
|
||||
future_observed_values: torch.Tensor,
|
||||
past_feat_dynamic_real: torch.Tensor,
|
||||
past_feat_dynamic_cat: torch.Tensor,
|
||||
feat_dynamic_real: torch.Tensor,
|
||||
feat_dynamic_cat: torch.Tensor,
|
||||
feat_static_real: torch.Tensor,
|
||||
feat_static_cat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
(
|
||||
past_covariates,
|
||||
future_covariates,
|
||||
static_covariates,
|
||||
offset,
|
||||
scale,
|
||||
) = self._preprocess(
|
||||
past_target,
|
||||
past_observed_values,
|
||||
past_feat_dynamic_real,
|
||||
past_feat_dynamic_cat,
|
||||
feat_dynamic_real,
|
||||
feat_dynamic_cat,
|
||||
feat_static_real,
|
||||
feat_static_cat,
|
||||
)
|
||||
|
||||
preds = super().forward(
|
||||
past_observed_values,
|
||||
past_covariates,
|
||||
future_covariates,
|
||||
static_covariates,
|
||||
)
|
||||
|
||||
preds = self._postprocess(preds, offset, scale)
|
||||
|
||||
loss = self.loss(future_target, preds)
|
||||
loss = weighted_average(loss, future_observed_values)
|
||||
return loss.mean()
|
||||
|
||||
|
||||
class TemporalFusionTransformerPredictionNetwork(TemporalFusionTransformerNetwork):
|
||||
def forward(
|
||||
self,
|
||||
past_target: torch.Tensor,
|
||||
past_observed_values: torch.Tensor,
|
||||
past_feat_dynamic_real: torch.Tensor,
|
||||
past_feat_dynamic_cat: torch.Tensor,
|
||||
feat_dynamic_real: torch.Tensor,
|
||||
feat_dynamic_cat: torch.Tensor,
|
||||
feat_static_real: torch.Tensor,
|
||||
feat_static_cat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
(
|
||||
past_covariates,
|
||||
future_covariates,
|
||||
static_covariates,
|
||||
offset,
|
||||
scale,
|
||||
) = self._preprocess(
|
||||
past_target,
|
||||
past_observed_values,
|
||||
past_feat_dynamic_real,
|
||||
past_feat_dynamic_cat,
|
||||
feat_dynamic_real,
|
||||
feat_dynamic_cat,
|
||||
feat_static_real,
|
||||
feat_static_cat,
|
||||
)
|
||||
|
||||
preds = super().forward(
|
||||
past_observed_values,
|
||||
past_covariates,
|
||||
future_covariates,
|
||||
static_covariates,
|
||||
)
|
||||
|
||||
preds = self._postprocess(preds, offset, scale)
|
||||
return preds.permute(0, 2, 1)
|
||||
@@ -0,0 +1,88 @@
|
||||
from typing import List, Optional, Tuple
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from gluonts.core.component import validated
|
||||
|
||||
|
||||
class QuantileLoss(nn.Module):
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
quantiles: List[float],
|
||||
quantile_weights: Optional[List[float]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.quantiles = quantiles
|
||||
self.num_quantiles = len(quantiles)
|
||||
self.quantile_weights = (
|
||||
[1.0 / self.num_quantiles for i in range(self.num_quantiles)]
|
||||
if not quantile_weights
|
||||
else quantile_weights
|
||||
)
|
||||
|
||||
def forward(self, y_true: torch.Tensor, y_pred: torch.Tensor, sample_weight=None):
|
||||
if self.num_quantiles > 1:
|
||||
y_pred_all = torch.chunk(y_pred, self.num_quantiles, dim=-1)
|
||||
else:
|
||||
y_pred_all = [y_pred]
|
||||
|
||||
qt_loss = []
|
||||
for i, y_pred_q in enumerate(y_pred_all):
|
||||
q = self.quantiles[i]
|
||||
weighted_qt = (
|
||||
self.compute_quantile_loss(y_true, y_pred_q.squeeze(-1), q)
|
||||
* self.quantile_weights[i]
|
||||
)
|
||||
qt_loss.append(weighted_qt)
|
||||
stacked_qt_losses = torch.stack(qt_loss, dim=-1)
|
||||
sum_qt_loss = torch.mean(stacked_qt_losses, dim=-1)
|
||||
if sample_weight is not None:
|
||||
return sample_weight * sum
|
||||
else:
|
||||
return sum_qt_loss
|
||||
|
||||
@staticmethod
|
||||
def compute_quantile_loss(
|
||||
y_true: torch.Tensor, y_pred_p: torch.Tensor, p: float
|
||||
) -> torch.Tensor:
|
||||
under_bias = p * torch.clamp(y_true - y_pred_p, min=0)
|
||||
over_bias = (1 - p) * torch.clamp(y_pred_p - y_true, min=0)
|
||||
|
||||
qt_loss = 2 * (under_bias + over_bias)
|
||||
return qt_loss
|
||||
|
||||
|
||||
class ProjectParams(nn.Module):
|
||||
@validated()
|
||||
def __init__(self, in_features, num_quantiles):
|
||||
super().__init__()
|
||||
self.projection = nn.Linear(in_features=in_features, out_features=num_quantiles)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.projection(x)
|
||||
|
||||
|
||||
class QuantileOutput:
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
input_size,
|
||||
quantiles: List[float],
|
||||
quantile_weights: Optional[List[float]] = None,
|
||||
) -> None:
|
||||
self.input_size = input_size
|
||||
self.quantiles = quantiles
|
||||
self.quantile_weights = quantile_weights
|
||||
|
||||
def get_loss(self) -> nn.Module:
|
||||
return QuantileLoss(
|
||||
quantiles=self.quantiles, quantile_weights=self.quantile_weights
|
||||
)
|
||||
|
||||
def get_quantile_proj(self) -> nn.Module:
|
||||
return ProjectParams(
|
||||
in_features=self.input_size, num_quantiles=len(self.quantiles)
|
||||
)
|
||||
@@ -0,0 +1,138 @@
|
||||
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License").
|
||||
# You may not use this file except in compliance with the License.
|
||||
# A copy of the License is located at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# or in the "license" file accompanying this file. This file is distributed
|
||||
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
|
||||
# express or implied. See the License for the specific language governing
|
||||
# permissions and limitations under the License.
|
||||
|
||||
from collections import Counter
|
||||
from typing import Iterator, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from gluonts.core.component import validated
|
||||
from gluonts.dataset.common import DataEntry
|
||||
from gluonts.dataset.field_names import FieldName
|
||||
from gluonts.transform import (
|
||||
InstanceSplitter,
|
||||
MapTransformation,
|
||||
shift_timestamp,
|
||||
target_transformation_length,
|
||||
)
|
||||
|
||||
|
||||
class BroadcastTo(MapTransformation):
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
field: str,
|
||||
ext_length: int = 0,
|
||||
target_field: str = FieldName.TARGET,
|
||||
) -> None:
|
||||
self.field = field
|
||||
self.ext_length = ext_length
|
||||
self.target_field = target_field
|
||||
|
||||
def map_transform(self, data: DataEntry, is_train: bool) -> DataEntry:
|
||||
length = target_transformation_length(
|
||||
data[self.target_field], self.ext_length, is_train
|
||||
)
|
||||
data[self.field] = np.broadcast_to(
|
||||
data[self.field],
|
||||
(data[self.field].shape[:-1] + (length,)),
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
class TFTInstanceSplitter(InstanceSplitter):
|
||||
@validated()
|
||||
def __init__(
|
||||
self,
|
||||
instance_sampler,
|
||||
past_length: int,
|
||||
future_length: int,
|
||||
target_field: str = FieldName.TARGET,
|
||||
is_pad_field: str = FieldName.IS_PAD,
|
||||
start_field: str = FieldName.START,
|
||||
forecast_start_field: str = FieldName.FORECAST_START,
|
||||
observed_value_field: str = FieldName.OBSERVED_VALUES,
|
||||
lead_time: int = 0,
|
||||
output_NTC: bool = True,
|
||||
time_series_fields: Optional[List[str]] = None,
|
||||
past_time_series_fields: Optional[List[str]] = None,
|
||||
dummy_value: float = 0.0,
|
||||
) -> None:
|
||||
|
||||
assert past_length > 0, "The value of `past_length` should be > 0"
|
||||
assert future_length > 0, "The value of `future_length` should be > 0"
|
||||
|
||||
self.instance_sampler = instance_sampler
|
||||
self.past_length = past_length
|
||||
self.future_length = future_length
|
||||
self.lead_time = lead_time
|
||||
self.output_NTC = output_NTC
|
||||
self.dummy_value = dummy_value
|
||||
|
||||
self.target_field = target_field
|
||||
self.is_pad_field = is_pad_field
|
||||
self.start_field = start_field
|
||||
self.forecast_start_field = forecast_start_field
|
||||
self.observed_value_field = observed_value_field
|
||||
|
||||
self.ts_fields = time_series_fields or []
|
||||
self.past_ts_fields = past_time_series_fields or []
|
||||
|
||||
def flatmap_transform(self, data: DataEntry, is_train: bool) -> Iterator[DataEntry]:
|
||||
pl = self.future_length
|
||||
lt = self.lead_time
|
||||
target = data[self.target_field]
|
||||
|
||||
sampled_indices = self.instance_sampler(target)
|
||||
|
||||
slice_cols = (
|
||||
self.ts_fields
|
||||
+ self.past_ts_fields
|
||||
+ [self.target_field, self.observed_value_field]
|
||||
)
|
||||
for i in sampled_indices:
|
||||
pad_length = max(self.past_length - i, 0)
|
||||
d = data.copy()
|
||||
|
||||
for field in slice_cols:
|
||||
if i >= self.past_length:
|
||||
past_piece = d[field][..., i - self.past_length : i]
|
||||
else:
|
||||
pad_block = np.full(
|
||||
shape=d[field].shape[:-1] + (pad_length,),
|
||||
fill_value=self.dummy_value,
|
||||
dtype=d[field].dtype,
|
||||
)
|
||||
past_piece = np.concatenate([pad_block, d[field][..., :i]], axis=-1)
|
||||
future_piece = d[field][..., (i + lt) : (i + lt + pl)]
|
||||
if field in self.ts_fields:
|
||||
piece = np.concatenate([past_piece, future_piece], axis=-1)
|
||||
if self.output_NTC:
|
||||
piece = piece.transpose()
|
||||
d[field] = piece
|
||||
else:
|
||||
if self.output_NTC:
|
||||
past_piece = past_piece.transpose()
|
||||
future_piece = future_piece.transpose()
|
||||
if field not in self.past_ts_fields:
|
||||
d[self._past(field)] = past_piece
|
||||
d[self._future(field)] = future_piece
|
||||
del d[field]
|
||||
else:
|
||||
d[field] = past_piece
|
||||
pad_indicator = np.zeros(self.past_length)
|
||||
if pad_length > 0:
|
||||
pad_indicator[:pad_length] = 1
|
||||
d[self._past(self.is_pad_field)] = pad_indicator
|
||||
d[self.forecast_start_field] = shift_timestamp(d[self.start_field], i + lt)
|
||||
yield d
|
||||
Reference in New Issue
Block a user