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

fix api to nn.Transformer

Browse Source
This commit is contained in:
Kashif Rasul
2020-01-27 19:57:54 +01:00
parent 2dc6f23efe
commit b1d1cdbea6
2 changed files with 84 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/transformer/transformer_estimator.py
+30 -28
View File
@@ -33,9 +33,11 @@ from pts.feature import (
from .transfomer_network import TransformerTrainingNetwork, TransformerPredictionNetwork
class TransformerEstimator(PTSEstimator):
def __init__(
self,
input_size: int,
freq: str,
prediction_length: int,
context_length: Optional[int] = None,
@@ -44,12 +46,11 @@ class TransformerEstimator(PTSEstimator):
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",
dim_feedforward: int = 256,
act_type: str = "gelu",
num_heads: int = 8,
num_encoder_layers: int 3,
num_decoder_layers: int 3,
scaling: bool = True,
lags_seq: Optional[List[int]] = None,
time_features: Optional[List[TimeFeature]] = None,
@@ -59,6 +60,7 @@ class TransformerEstimator(PTSEstimator):
) -> None:
super().__init__(trainer=trainer)
self.input_size = input_size
self.freq = freq
self.prediction_length = prediction_length
self.context_length = (
@@ -72,7 +74,8 @@ class TransformerEstimator(PTSEstimator):
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)
lags_seq if lags_seq is not None else get_lags_for_frequency(
freq_str=freq)
)
self.time_features = (
time_features
@@ -82,22 +85,11 @@ class TransformerEstimator(PTSEstimator):
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_"
# )
self.num_heads = num_heads
self.act_type = act_type
self.dim_feedforward = dim_feedforward
self.num_encoder_layers = num_encoder_layers
self.num_decoder_layers = num_decoder_layers
def create_transformation(self) -> Transformation:
remove_field_names = [
@@ -166,8 +158,13 @@ class TransformerEstimator(PTSEstimator):
def create_training_network(self, device: torch.device) -> TransformerTrainingNetwork:
training_network = TransformerTrainingNetwork(
encoder=self.encoder,
decoder=self.decoder,
input_size=self.input_size,
num_heads=self.num_heads,
act_type=self.act_type,
dropout_rate=self.dropout_rate,
dim_feedforward=self.dim_feedforward,
num_encoder_layers=self.num_encoder_layers,
num_decoder_layers=self.num_decoder_layers,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
@@ -175,7 +172,7 @@ class TransformerEstimator(PTSEstimator):
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=True,
scaling=self.scaling,
).to(device)
return training_network
@@ -185,8 +182,13 @@ class TransformerEstimator(PTSEstimator):
) -> Predictor:
prediction_network = TransformerPredictionNetwork(
encoder=self.encoder,
decoder=self.decoder,
input_size=self.input_size,
num_heads=self.num_heads,
act_type=self.act_type,
dropout_rate=self.dropout_rate,
dim_feedforward=self.dim_feedforward,
num_encoder_layers=self.num_encoder_layers,
num_decoder_layers=self.num_decoder_layers,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
@@ -194,7 +196,7 @@ class TransformerEstimator(PTSEstimator):
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
lags_seq=self.lags_seq,
scaling=True,
scaling=self.scaling,
num_parallel_samples=self.num_parallel_samples,
)
pts/model/transformer/transformer_network.py
+54 -45
View File
@@ -17,8 +17,13 @@ from .trans_decoder import TransformerDecoder
class TransformerNetwork(nn.Module):
def __init__(
self,
encoder: TransformerEncoder,
decoder: TransformerDecoder,
input_size: int,
num_heads: int,
act_type: str,
dropout_rate: float,
dim_feedforward: int,
num_encoder_layers: int,
num_decoder_layers: int,
history_length: int,
context_length: int,
prediction_length: int,
@@ -47,19 +52,19 @@ class TransformerNetwork(nn.Module):
self.lags_seq = lags_seq
self.target_shape = distr_output.event_shape
self.transformer = nn.Transformer(
d_model=input_size,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6,
dim_feedforward=2048,
dropout=0.1,
activation='relu',
d_model=input_size,
nhead=num_heads,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward,
dropout=dropout_rate,
activation=act_type,
)
self.proj_dist_args = distr_output.get_args_proj(input_size)
self.embedder = FeatureEmbedder(
cardinalities=cardinality,
embedding_dims=[embedding_dimension for _ in cardinality],
@@ -70,7 +75,6 @@ class TransformerNetwork(nn.Module):
else:
self.scaler = NOPScaler(keepdims=True)
@staticmethod
def get_lagged_subsequences(
sequence: torch.Tensor,
@@ -114,13 +118,15 @@ class TransformerNetwork(nn.Module):
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)
# (batch_size, num_features, history_length)
past_time_feat: torch.Tensor,
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)
# (batch_size, prediction_length)
future_target: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Creates inputs for the transformer network.
@@ -128,29 +134,32 @@ class TransformerNetwork(nn.Module):
"""
if future_time_feat is None or future_target is None:
time_feat = past_time_feat[:, self.history_length - self.context_length:, ...] #.slice_axis(
# .slice_axis(
time_feat = past_time_feat[:,
self.history_length - self.context_length:, ...]
# 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
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)
# (batch_size, sub_seq_len, *target_shape, num_lags)
lags = self.get_lagged_subsequences(
sequence=sequence,
sequence_length=sequence_length,
@@ -161,10 +170,10 @@ class TransformerNetwork(nn.Module):
# 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(
past_target[:, -self.context_length:, ...], # .slice_axis(
# axis=1, begin=-self.context_length, end=None
# ),
past_observed_values[:,-self.context_length:,...] #.slice_axis(
past_observed_values[:, -self.context_length:, ...] # .slice_axis(
# axis=1, begin=-self.context_length, end=None
# ),
)
@@ -194,11 +203,11 @@ class TransformerNetwork(nn.Module):
)
# (batch_size, sub_seq_len, input_dim)
inputs = torch.cat((input_lags, time_feat, repeated_static_feat), dim=-1)
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))
@@ -206,7 +215,6 @@ class TransformerNetwork(nn.Module):
mask = mask + torch.eye(d, d, k + 1)
return mask
class TransformerTrainingNetwork(TransformerNetwork):
# noinspection PyMethodOverriding,PyPep8Naming
@@ -244,10 +252,10 @@ class TransformerTrainingNetwork(TransformerNetwork):
future_target=future_target,
)
enc_input = input[:, :self.context_length, ...] # F.slice_axis(
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(
dec_input = input[:, self.context_length:, ...] # F.slice_axis(
# inputs, axis=1, begin=self.context_length, end=None
# )
@@ -257,8 +265,8 @@ class TransformerTrainingNetwork(TransformerNetwork):
# input to decoder
dec_output = self.transformer.decoder(
dec_input,
enc_out, #memory
self.upper_triangular_mask(self.prediction_length), #mask
enc_out, # memory
self.upper_triangular_mask(self.prediction_length), # mask
)
# compute loss
@@ -267,7 +275,7 @@ class TransformerTrainingNetwork(TransformerNetwork):
loss = distr.loss(future_target)
return loss.mean()
class TransformerPredictionNetwork(TransformerNetwork):
def __init__(self, num_parallel_samples: int = 100, **kwargs) -> None:
@@ -353,13 +361,14 @@ class TransformerPredictionNetwork(TransformerNetwork):
# (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_time_feat[:, k:k+1, :],
repeated_static_feat),
dim=-1,
)
dec_output = self.transformer.decoder(dec_input, repeated_enc_out, None)
dec_output = self.transformer.decoder(
dec_input, repeated_enc_out, None)
distr_args = self.proj_dist_args(dec_output)
# compute likelihood of target given the predicted parameters