wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-19 11:27:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

initial transformer model

Browse Source 
todo decoder encoder
This commit is contained in:
Dr. Kashif Rasul
2020-01-27 16:24:02 +01:00
parent d1005e2aa5
commit 3184f5d588
5 changed files with 642 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/transformer/__init__.py
+1
View File
@@ -0,0 +1 @@
from .transformer_estimator import TransformerEstimator
pts/model/transformer/transformer_decoder.py
View File
pts/model/transformer/transformer_encoder.py
View File
pts/model/transformer/transformer_estimator.py
+213
View File
@@ -0,0 +1,213 @@
from typing import List, Optional
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from pts import Trainer
from pts.model import PTSEstimator, PTSPredictor, copy_parameters
from pts.modules import RealNVP
from pts.dataset import FieldName
from pts.transform import (
Transformation,
Chain,
InstanceSplitter,
ExpectedNumInstanceSampler,
CDFtoGaussianTransform,
cdf_to_gaussian_forward_transform,
RenameFields,
AsNumpyArray,
ExpandDimArray,
AddObservedValuesIndicator,
AddTimeFeatures,
VstackFeatures,
SetFieldIfNotPresent,
TargetDimIndicator,
)
from pts.feature import (
TimeFeature,
fourier_time_features_from_frequency_str,
get_fourier_lags_for_frequency,
)
from .transformer_decoder import TransformerDecoder
from .transfomer_encoder import TransformerEncoder
from .transfomer_network import TransformerTrainingNetwork, TransformerPredictionNetwork
class TransformerEstimator(PTSEstimator):
def __init__(
self,
freq: str,
prediction_length: int,
context_length: Optional[int] = None,
trainer: Trainer = Trainer(),
dropout_rate: float = 0.1,
cardinality: Optional[List[int]] = None,
embedding_dimension: int = 20,
distr_output: DistributionOutput = StudentTOutput(),
model_dim: int = 32,
inner_ff_dim_scale: int = 4,
pre_seq: str = "dn",
post_seq: str = "drn",
act_type: str = "softrelu",
num_heads: int = 8,
scaling: bool = True,
lags_seq: Optional[List[int]] = None,
time_features: Optional[List[TimeFeature]] = None,
use_feat_dynamic_real: bool = False,
use_feat_static_cat: bool = False,
num_parallel_samples: int = 100,
) -> None:
super().__init__(trainer=trainer)
self.freq = freq
self.prediction_length = prediction_length
self.context_length = (
context_length if context_length is not None else prediction_length
)
self.distr_output = distr_output
self.dropout_rate = dropout_rate
self.use_feat_dynamic_real = use_feat_dynamic_real
self.use_feat_static_cat = use_feat_static_cat
self.cardinality = cardinality if use_feat_static_cat else [1]
self.embedding_dimension = embedding_dimension
self.num_parallel_samples = num_parallel_samples
self.lags_seq = (
lags_seq if lags_seq is not None else get_lags_for_frequency(freq_str=freq)
)
self.time_features = (
time_features
if time_features is not None
else time_features_from_frequency_str(self.freq)
)
self.history_length = self.context_length + max(self.lags_seq)
self.scaling = scaling
self.config = {
"model_dim": model_dim,
"pre_seq": pre_seq,
"post_seq": post_seq,
"dropout_rate": dropout_rate,
"inner_ff_dim_scale": inner_ff_dim_scale,
"act_type": act_type,
"num_heads": num_heads,
}
self.encoder = TransformerEncoder(
self.context_length, self.config, prefix="enc_"
)
self.decoder = TransformerDecoder(
self.prediction_length, self.config, prefix="dec_"
)
def create_transformation(self) -> Transformation:
remove_field_names = [
FieldName.FEAT_DYNAMIC_CAT,
FieldName.FEAT_STATIC_REAL,
]
if not self.use_feat_dynamic_real:
remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)
return Chain(
[RemoveFields(field_names=remove_field_names)]
+ (
[SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0.0])]
if not self.use_feat_static_cat
else []
)
+ [
AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1),
AsNumpyArray(
field=FieldName.TARGET,
# in the following line, we add 1 for the time dimension
expected_ndim=1 + len(self.distr_output.event_shape),
),
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,
),
VstackFeatures(
output_field=FieldName.FEAT_TIME,
input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
+ (
[FieldName.FEAT_DYNAMIC_REAL]
if self.use_feat_dynamic_real
else []
),
),
InstanceSplitter(
target_field=FieldName.TARGET,
is_pad_field=FieldName.IS_PAD,
start_field=FieldName.START,
forecast_start_field=FieldName.FORECAST_START,
train_sampler=ExpectedNumInstanceSampler(num_instances=1),
past_length=self.history_length,
future_length=self.prediction_length,
time_series_fields=[
FieldName.FEAT_TIME,
FieldName.OBSERVED_VALUES,
],
),
]
)
def create_training_network(self, device: torch.device) -> TransformerTrainingNetwork:
training_network = TransformerTrainingNetwork(
encoder=self.encoder,
decoder=self.decoder,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
distr_output=self.distr_output,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=True,
).to(device)
return training_network
def create_predictor(
self, transformation: Transformation, trained_network: nn.Module, device: torch.device,
) -> Predictor:
prediction_network = TransformerPredictionNetwork(
encoder=self.encoder,
decoder=self.decoder,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
distr_output=self.distr_output,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=True,
num_parallel_samples=self.num_parallel_samples,
)
copy_parameters(trained_network, prediction_network)
return PTSPredictor(
input_transform=transformation,
prediction_net=prediction_network,
batch_size=self.trainer.batch_size,
freq=self.freq,
prediction_length=self.prediction_length,
device=device,
output_transform=None,
)
pts/model/transformer/transformer_network.py
+428
View File
@@ -0,0 +1,428 @@
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
from torch.distributions import Distribution
import numpy as np
from pts.modules import DistributionOutput, MeanScaler, NOPScaler, FeatureEmbedder
from pts.model import weighted_average
from .trans_encoder import TransformerEncoder
from .trans_decoder import TransformerDecoder
LARGE_NEGATIVE_VALUE = -99999999
class TransformerNetwork(nn.Module):
def __init__(
self,
encoder: TransformerEncoder,
decoder: TransformerDecoder,
history_length: int,
context_length: int,
prediction_length: int,
distr_output: DistributionOutput,
cardinality: List[int],
embedding_dimension: int,
lags_seq: List[int],
scaling: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.history_length = history_length
self.context_length = context_length
self.prediction_length = prediction_length
self.scaling = scaling
self.cardinality = cardinality
self.embedding_dimension = embedding_dimension
self.distr_output = distr_output
assert len(set(lags_seq)) == len(
lags_seq
), "no duplicated lags allowed!"
lags_seq.sort()
self.lags_seq = lags_seq
self.target_shape = distr_output.event_shape
self.proj_dist_args = distr_output.get_args_proj(input_size) #TODO figure it out
self.encoder = encoder
self.decoder = decoder
self.embedder = FeatureEmbedder(
cardinalities=cardinality,
embedding_dims=[embedding_dimension for _ in cardinality],
)
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 : 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 create_network_input(
self,
feat_static_cat: torch.Tensor, # (batch_size, num_features)
past_time_feat: torch.Tensor, # (batch_size, num_features, history_length)
past_target: torch.Tensor, # (batch_size, history_length, 1)
past_observed_values: torch.Tensor, # (batch_size, history_length)
future_time_feat: Optional[
torch.Tensor
], # (batch_size, num_features, prediction_length)
future_target: Optional[torch.Tensor], # (batch_size, prediction_length)
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Creates inputs for the transformer network.
All tensor arguments should have NTC layout.
"""
if future_time_feat is None or future_target is None:
time_feat = past_time_feat[:, self.history_length - self.context_length:, ...] #.slice_axis(
# axis=1,
# begin=self.history_length - self.context_length,
# end=None,
)
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:,...], #.slice_axis(
# axis=1,
# begin=self.history_length - self.context_length,
# end=None,
# ),
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
# (batch_size, sub_seq_len, *target_shape, 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_shape)
_, scale = self.scaler(
past_target[:,-self.context_length:,...], #.slice_axis(
# axis=1, begin=-self.context_length, end=None
# ),
past_observed_values[:,-self.context_length:,...] #.slice_axis(
# axis=1, begin=-self.context_length, end=None
# ),
)
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,
torch.log(scale)
if len(self.target_shape) == 0
else torch.log(scale.squeeze(1))),
dim=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)
return inputs, scale, static_feat
@staticmethod
def upper_triangular_mask(d):
mask = torch.zeros_like(torch.eye(d))
for k in range(d - 1):
mask = mask + torch.eye(d, d, k + 1)
return mask * LARGE_NEGATIVE_VALUE
class TransformerTrainingNetwork(TransformerNetwork):
# noinspection PyMethodOverriding,PyPep8Naming
def forward(
self,
feat_static_cat: 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,
) -> torch.Tensor:
"""
Computes the loss for training Transformer, all inputs tensors representing time series have NTC layout.
Parameters
----------
F
feat_static_cat : (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, seq_len)
future_time_feat : (batch_size, prediction_length, num_features)
future_target : (batch_size, prediction_length, *target_shape)
Returns
-------
Loss with shape (batch_size, context + prediction_length, 1)
"""
# create the inputs for the encoder
inputs, scale, _ = self.create_network_input(
feat_static_cat=feat_static_cat,
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,
)
enc_input = input[:, :self.context_length, ...] # F.slice_axis(
# inputs, axis=1, begin=0, end=self.context_length
# )
dec_input = input[:,self.context_length:,...] #F.slice_axis(
# inputs, axis=1, begin=self.context_length, end=None
# )
# pass through encoder
enc_out = self.encoder(enc_input)
# input to decoder
dec_output = self.decoder(
dec_input,
enc_out,
self.upper_triangular_mask(F, self.prediction_length),
)
# compute loss
distr_args = self.proj_dist_args(dec_output)
distr = self.distr_output.distribution(distr_args, scale=scale)
loss = distr.loss(future_target)
return loss.mean()
class TransformerPredictionNetwork(TransformerNetwork):
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,
enc_out: torch.Tensor,
) -> torch.Tensor:
"""
Computes sample paths by decoding from the transformer.
Parameters
----------
static_feat : Tensor
static features. Shape: (batch_size, num_static_features).
past_target : Tensor
target history. Shape: (batch_size, history_length, 1).
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, ).
enc_out: Tensor
output of the encoder. Shape: (batch_size, num_cells)
Returns
--------
sample_paths : 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_enc_out = enc_out.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
).unsqueeze(1)
repeated_scale = scale.repeat_interleave(
repeats=self.num_parallel_samples, dim=0
)
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,
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)
# 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 = lags_scaled.reshape(
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)
dec_input = torch.cat((
input_lags,
repeated_time_feat[:,k:k+1,:],
repeated_static_feat),
dim=-1,
)
dec_output = self.decoder(dec_input, repeated_enc_out, None, False)
distr_args = self.proj_dist_args(dec_output)
# 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)
# reset cache of the decoder
self.decoder.cache_reset()
# (batch_size * num_samples, prediction_length, *target_shape)
samples = torch.cat(future_samples, dim=1)
# (batch_size, num_samples, *target_shape, prediction_length)
return samples.reshape(
(
(-1, self.num_parallel_samples)
+ self.target_shape
+ (self.prediction_length,)
)
)
def forward(
self,
feat_static_cat: torch.Tensor,
past_time_feat: torch.Tensor,
past_target: torch.Tensor,
past_observed_values: torch.Tensor,
future_time_feat: torch.Tensor,
) -> torch.Tensor:
"""
Predicts samples, all tensors should have NTC layout.
Parameters
----------
feat_static_cat : (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 predicted samples
-------
"""
# create the inputs for the encoder
inputs, scale, static_feat = self.create_network_input(
feat_static_cat=feat_static_cat,
past_time_feat=past_time_feat,
past_target=past_target,
past_observed_values=past_observed_values,
future_time_feat=None,
future_target=None,
)
# pass through encoder
enc_out = self.encoder(inputs)
return self.sampling_decoder(
past_target=past_target,
time_feat=future_time_feat,
static_feat=static_feat,
scale=scale,
enc_out=enc_out,
)