From 3e3f46bec49e99c351fa3a0a2de56ede561e5e5c Mon Sep 17 00:00:00 2001 From: Hstellar Date: Tue, 12 Apr 2022 13:28:37 -0400 Subject: [PATCH] Adding pyraformer(not working) --- Pyraformer/.DS_Store | Bin 0 -> 6148 bytes Pyraformer/.typesafe | 0 Pyraformer/__init__.py | 11 + Pyraformer/estimator.py | 343 +++++++++ Pyraformer/lightning_module.py | 82 +++ Pyraformer/module.py | 705 +++++++++++++++++++ Pyraformer/pyraformer.ipynb | 1 + Pyraformer/pyraformer/Layers.py | 408 +++++++++++ Pyraformer/pyraformer/Modules.py | 25 + Pyraformer/pyraformer/PAM_TVM.py | 65 ++ Pyraformer/pyraformer/SubLayers.py | 96 +++ Pyraformer/pyraformer/embed.py | 158 +++++ Pyraformer/pyraformer/graph_attention.py | 580 +++++++++++++++ Pyraformer/pyraformer/hierarchical_mm_tvm.py | 282 ++++++++ Pyraformer/tools.py | 93 +++ 15 files changed, 2849 insertions(+) create mode 100644 Pyraformer/.DS_Store create mode 100644 Pyraformer/.typesafe create mode 100644 Pyraformer/__init__.py create mode 100644 Pyraformer/estimator.py create mode 100644 Pyraformer/lightning_module.py create mode 100644 Pyraformer/module.py create mode 100644 Pyraformer/pyraformer.ipynb create mode 100644 Pyraformer/pyraformer/Layers.py create mode 100644 Pyraformer/pyraformer/Modules.py create mode 100644 Pyraformer/pyraformer/PAM_TVM.py create mode 100644 Pyraformer/pyraformer/SubLayers.py create mode 100644 Pyraformer/pyraformer/embed.py create mode 100644 Pyraformer/pyraformer/graph_attention.py create mode 100644 Pyraformer/pyraformer/hierarchical_mm_tvm.py create mode 100644 Pyraformer/tools.py diff --git a/Pyraformer/.DS_Store b/Pyraformer/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..38734ca2de71d90578b12a191d5ff30a57f26d5c GIT binary patch literal 6148 zcmeHKJ8Hu~5S@u#2;8`IxmU;y7U7)02atb(6aoc8igc=cE+5TrJ{W}TCXgn)ftj~E znx|!7q0xwlw%_Mhk+q04a6`FRn43K}pV>=h6bQ#VPI7>M$h&m2>c}_IdtuEgkdX>d0V;4;z`hR!Zden?K>u_g_y_=8 zBJGB?&l13531Cee1CfDgP=P_!95FQL$d|0EiDO{UMRWMjJXv!>Q9m8;FJ3NM0~x6R z6__fpi0#_?{~P?t{68geM+KA!1r*gxx&q`b_#;GW1zQV hY^)t`yeR65t?|4jj)6``-swR8445u7D)83|+yF)|6_)@2 literal 0 HcmV?d00001 diff --git a/Pyraformer/.typesafe b/Pyraformer/.typesafe new file mode 100644 index 0000000..e69de29 diff --git a/Pyraformer/__init__.py b/Pyraformer/__init__.py new file mode 100644 index 0000000..03d1301 --- /dev/null +++ b/Pyraformer/__init__.py @@ -0,0 +1,11 @@ +from .estimator import PyraformerEstimator +from .lightning_module import PyraformerLightningModule +from .module import PyraformerSSModel +from .module import PyraformerLRModel + +__all__ = [ + "PyraformerSSModel", + "PyraformerLRModel", + "PyraformerLightningModule", + "PyraformerEstimator", +] diff --git a/Pyraformer/estimator.py b/Pyraformer/estimator.py new file mode 100644 index 0000000..b90232d --- /dev/null +++ b/Pyraformer/estimator.py @@ -0,0 +1,343 @@ +from typing import Any, Dict, Iterable, List, Optional + +import torch +from gluonts.core.component import validated +from gluonts.dataset.common import Dataset +from gluonts.dataset.field_names import FieldName +from gluonts.itertools import Cyclic, IterableSlice, PseudoShuffled +from gluonts.time_feature import TimeFeature, time_features_from_frequency_str +from gluonts.torch.model.estimator import PyTorchLightningEstimator +from gluonts.torch.model.predictor import PyTorchPredictor +from gluonts.torch.modules.distribution_output import DistributionOutput, StudentTOutput +from gluonts.torch.modules.loss import DistributionLoss, NegativeLogLikelihood +from gluonts.time_feature import get_lags_for_frequency +from gluonts.torch.util import IterableDataset +from gluonts.transform import ( + AddAgeFeature, + AddObservedValuesIndicator, + AddTimeFeatures, + AsNumpyArray, + Chain, + ExpectedNumInstanceSampler, + InstanceSplitter, + RemoveFields, + SelectFields, + SetField, + TestSplitSampler, + Transformation, + ValidationSplitSampler, + VstackFeatures, +) +from gluonts.transform.sampler import InstanceSampler +from lightning_module import PyraformerLightningModule +from module import PyraformerSSModel +from module import PyraformerLRModel +from torch.utils.data import DataLoader +from tools import SingleStepLoss as LossFactory +from torch.utils.data.sampler import RandomSampler +PREDICTION_INPUT_NAMES = [ + "feat_static_cat", + "feat_static_real", + "past_time_feat", + "past_target", + "past_observed_values", + "future_time_feat", +] + +TRAINING_INPUT_NAMES = PREDICTION_INPUT_NAMES + [ + "future_target", + "future_observed_values", +] + + +class PyraformerEstimator(PyTorchLightningEstimator): + @validated() + def __init__( + self, + freq: str, + prediction_length: int, + #Train parameters + inner_batch: int = 8, + lr: float = 1e-5, + visualize_fre: int = 2000, + pretrain: bool = True, + hard_sample_mining:bool=True, + covariate_size: int = 3, + + # Model parameters + num_seq: int = 370,# + decoder: str = 'FC',# selection: [FC, attention] + context_length: Optional[int] = None, + input_size: int = 1, + dropout: float = 0.1, + d_model: int = 512, + d_inner_hid: int = 512, + d_k: int = 128, + d_v:int = 128, + num_heads: int = 4, + n_layer: int = 4, + # loss: DistributionLoss = LossFactory, + ignore_zero: bool = True, + single_step: bool = True,#if False, Multistep=True + + inner_size: int = 3, + use_tvm: bool = False, + num_feat_dynamic_real: int = 0, + num_feat_static_cat: int = 0, + num_feat_static_real: int = 0, + cardinality: Optional[List[int]] = None, + embedding_dimension: Optional[List[int]] = None, + distr_output: DistributionOutput = StudentTOutput(), + # loss: DistributionLoss = NegativeLogLikelihood(), + scaling: bool = True, + lags_seq: Optional[List[int]] = None, + time_features: Optional[List[TimeFeature]] = None, + num_parallel_samples: int = 100, + batch_size: int = 32, + num_batches_per_epoch: int = 50, + trainer_kwargs: Optional[Dict[str, Any]] = dict(), + train_sampler: Optional[InstanceSampler] = None, + validation_sampler: Optional[InstanceSampler] = None, + window_size: int = [4, 4, 4] + ) -> None: + trainer_kwargs = { + "max_epochs": 10, + **trainer_kwargs, + } + super().__init__(trainer_kwargs=trainer_kwargs) + + self.freq = freq + self.context_length = ( + context_length if context_length is not None else prediction_length + ) + self.inner_batch = inner_batch + self.lr = lr + # self.visualize_fre = visualize_fre + self.covariate_size = covariate_size + self.num_seq = num_seq + self.input_size = input_size + self.dropout = dropout + self.d_model = d_model + self.d_inner_hid = d_inner_hid + self.d_k = d_k + self.d_v = d_v + self.num_heads = num_heads + self.n_layer = n_layer + self.single_step = single_step + self.ignore_zero = ignore_zero + self.loss = LossFactory(self.ignore_zero) if self.single_step==True else torch.nn.MSELoss(reduction='none') + self.batch_size = batch_size + self.distr_output = distr_output + + self.window_size = window_size#[4,4,4]#window_size + self.inner_size = inner_size + self.use_tvm = use_tvm + self.prediction_length = prediction_length + # self.epochs = trainer_kwargs['max_epochs'] + # self.train_sampler = RandomSampler or ExpectedNumInstanceSampler(num_instances=1.0, min_future=prediction_length) + # self.validation_sampler = RandomSampler or ValidationSplitSampler(min_future=prediction_length) + # self.test_sampler = RandomSampler + + self.num_feat_dynamic_real = num_feat_dynamic_real + self.num_feat_static_cat = num_feat_static_cat + self.num_feat_static_real = num_feat_static_real + self.cardinality = ( + cardinality if cardinality and num_feat_static_cat > 0 else [1] + ) + self.embedding_dimension = embedding_dimension + self.scaling = scaling + self.lags_seq = lags_seq + self.time_features = ( + time_features + if time_features is not None + else time_features_from_frequency_str(self.freq) + ) + + self.num_parallel_samples = num_parallel_samples + self.batch_size = batch_size + self.num_batches_per_epoch = num_batches_per_epoch + + self.train_sampler = train_sampler or ExpectedNumInstanceSampler( + num_instances=1.0, min_future=prediction_length + ) + self.validation_sampler = validation_sampler or ValidationSplitSampler( + min_future=prediction_length + ) + + def create_transformation(self) -> Transformation: + remove_field_names = [] + if self.num_feat_static_real == 0: + remove_field_names.append(FieldName.FEAT_STATIC_REAL) + if self.num_feat_dynamic_real == 0: + 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])] + if not self.num_feat_static_cat > 0 + else [] + ) + + ( + [SetField(output_field=FieldName.FEAT_STATIC_REAL, value=[0.0])] + if not self.num_feat_static_real > 0 + else [] + ) + + [ + AsNumpyArray( + field=FieldName.FEAT_STATIC_CAT, + expected_ndim=1, + dtype=int, + ), + AsNumpyArray( + field=FieldName.FEAT_STATIC_REAL, + 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.num_feat_dynamic_real > 0 + else [] + ), + ), + ] + ) + + def _create_instance_splitter(self, module: PyraformerLightningModule, mode: str): + assert mode in ["training", "validation", "test"] + + instance_sampler = { + "training": self.train_sampler, + "validation": self.validation_sampler, + "test": TestSplitSampler(), + }[mode] + print(instance_sampler) + return InstanceSplitter( + target_field=FieldName.TARGET, + is_pad_field=FieldName.IS_PAD, + start_field=FieldName.START, + forecast_start_field=FieldName.FORECAST_START, + instance_sampler=instance_sampler, + past_length=module.model._past_length, + future_length=self.prediction_length, + time_series_fields=[ + FieldName.FEAT_TIME, + FieldName.OBSERVED_VALUES, + ], + dummy_value=self.distr_output.value_in_support, + ) + + def create_training_data_loader( + self, + data: Dataset, + module: PyraformerLightningModule, + shuffle_buffer_length: Optional[int] = None, + **kwargs, + ) -> Iterable: + transformation = self._create_instance_splitter( + module, "training" + ) + SelectFields(TRAINING_INPUT_NAMES) + + training_instances = transformation.apply( + Cyclic(data) + if shuffle_buffer_length is None + else PseudoShuffled( + Cyclic(data), shuffle_buffer_length=shuffle_buffer_length + ) + ) + + return IterableSlice( + iter( + DataLoader( + IterableDataset(training_instances), + batch_size=self.batch_size, + **kwargs, + ) + ), + self.num_batches_per_epoch, + ) + + def create_validation_data_loader( + self, + data: Dataset, + module: PyraformerLightningModule, + **kwargs, + ) -> Iterable: + transformation = self._create_instance_splitter( + module, "validation" + ) + SelectFields(TRAINING_INPUT_NAMES) + + validation_instances = transformation.apply(data) + + return DataLoader( + IterableDataset(validation_instances), + batch_size=self.batch_size, + **kwargs, + ) + + def create_predictor( + self, + transformation: Transformation, + module: PyraformerLightningModule, + ) -> PyTorchPredictor: + prediction_splitter = self._create_instance_splitter(module, "test") + + return PyTorchPredictor( + input_transform=transformation + prediction_splitter, + input_names=PREDICTION_INPUT_NAMES, + prediction_net=module.model, + batch_size=self.batch_size, + freq=self.freq, + prediction_length=self.prediction_length, + device=torch.device("cuda" if torch.cuda.is_available() else "cpu"), + ) + + def create_lightning_module(self) -> PyraformerLightningModule: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + if self.single_step: + model = PyraformerSSModel(freq= self.freq, covariate_size = self.covariate_size, + num_seq=self.num_seq, input_size = self.input_size, dropout = self.dropout, d_model = self.d_model, + d_inner_hid = self.d_inner_hid, d_k = self.d_k, d_v = self.d_v, + num_heads = self.num_heads, n_layer = self.n_layer, loss = self.loss, + window_size = self.window_size, inner_size = self.inner_size, + use_tvm = self.use_tvm, prediction_length = self.prediction_length,context_length = self.context_length, lags_seq = self.lags_seq,embedding_dimension=self.embedding_dimension, num_feat_dynamic_real= self.num_feat_dynamic_real, num_feat_dynamic_real, + num_feat_static_cat = self.num_feat_static_cat, + num_feat_static_real = self.num_feat_static_real, + cardinality = self.cardinality, + embedding_dimension = self.embedding_dimension, + distr_output=self.distr_output, + scaling=self.scaling,num_parallel_samples=self.num_parallel_samples, device=device) + # else: + # model = PyraformerLRModel(freq= self.freq, covariate_size = self.covariate_size, + # num_seq=self.num_seq, input_size = self.input_size, dropout = self.dropout, d_model = self.d_model, + # d_inner_hid = self.d_inner_hid, d_k = self.d_k, d_v = self.d_v, + # num_heads = self.num_heads, n_layer = self.n_layer, loss = self.loss, + # window_size = self.window_size, inner_size = self.inner_size, + # use_tvm = self.use_tvm, prediction_length = self.prediction_length,context_length = self.context_length, lags_seq = self.lags_seq, device=device) + return PyraformerLightningModule(model=model, loss=self.loss) + + diff --git a/Pyraformer/lightning_module.py b/Pyraformer/lightning_module.py new file mode 100644 index 0000000..5753cca --- /dev/null +++ b/Pyraformer/lightning_module.py @@ -0,0 +1,82 @@ +import pytorch_lightning as pl +import torch +from gluonts.torch.modules.loss import DistributionLoss, NegativeLogLikelihood +from gluonts.torch.util import weighted_average +from module import PyraformerSSModel +from module import PyraformerLRModel +from tools import SingleStepLoss as LossFactory +from tools import AE_loss +# from module import PyraformerModel + + +class PyraformerLightningModule(pl.LightningModule): + def __init__(self, model: PyraformerSSModel, loss: DistributionLoss = LossFactory, lr: float = 1e-5, weight_decay: float = 1e-8,) -> None: + super().__init__() + self.save_hyperparameters() + self.model = model + self.loss = loss + self.lr = lr + self.weight_decay = weight_decay + + def training_step(self, batch, batch_idx: int): + + """Execute training step""" + train_loss = self(batch) + self.log( + "train_loss", + train_loss, + on_epoch=True, + on_step=False, + prog_bar=True, + ) + return train_loss + + def validation_step(self, batch, batch_idx: int): + """Execute validation step""" + with torch.inference_mode(): + val_loss = self(batch) + self.log("val_loss", val_loss, on_epoch=True, on_step=False, prog_bar=True) + return val_loss + + def configure_optimizers(self): + """Returns the optimizer to use""" + return torch.optim.Adam( + self.model.parameters(), + lr=self.lr, + weight_decay=self.weight_decay, + ) + + def forward(self, batch): + feat_static_cat = batch["feat_static_cat"] + feat_static_real = batch["feat_static_real"] + past_time_feat = batch["past_time_feat"] + past_target = batch["past_target"] + future_time_feat = batch["future_time_feat"] + future_target = batch["future_target"] + past_observed_values = batch["past_observed_values"] + future_observed_values = batch["future_observed_values"] + + Pyraformer_inputs, scale, _ = self.model.create_network_inputs( + feat_static_cat, + feat_static_real, + past_time_feat, + past_target, + past_observed_values, + future_time_feat, + future_target, + ) + params = self.model.output_params(Pyraformer_inputs) + distr = self.model.output_distribution(params, scale) + + loss_values = self.loss(distr, future_target) + + if len(self.model.target_shape) == 0: + loss_weights = future_observed_values + else: + loss_weights = future_observed_values.min(dim=-1, keepdim=False) + + return weighted_average(loss_values, weights=loss_weights) + + + + diff --git a/Pyraformer/module.py b/Pyraformer/module.py new file mode 100644 index 0000000..d8a13d6 --- /dev/null +++ b/Pyraformer/module.py @@ -0,0 +1,705 @@ +from typing import List, Optional + +import torch +import torch.nn as nn +from gluonts.core.component import validated +from gluonts.time_feature import get_lags_for_frequency +from gluonts.torch.modules.distribution_output import DistributionOutput, StudentTOutput +from gluonts.torch.modules.feature import FeatureEmbedder +from gluonts.torch.modules.scaler import MeanScaler, NOPScaler + +from pyraformer.Layers import EncoderLayer, Predictor, Decoder +from pyraformer.Layers import Bottleneck_Construct, Conv_Construct, MaxPooling_Construct, AvgPooling_Construct +from pyraformer.Layers import get_mask, refer_points, get_k_q, get_q_k, get_subsequent_mask +from pyraformer.embed import SingleStepEmbedding, DataEmbedding, CustomEmbedding + + +class EncoderSS(nn.Module): + """ A encoder model with self attention mechanism. """ + def __init__(self, covariate_size, + num_seq, input_size ,dropout , d_model, + d_inner_hid, d_k, d_v , + num_heads , n_layer, loss , + window_size , inner_size, + use_tvm, prediction_length, device): + super().__init__() + + self.d_model = d_model + self.window_size = window_size + self.num_heads = num_heads + self.mask, self.all_size = get_mask(input_size, window_size, inner_size, device) + self.indexes = refer_points(self.all_size, window_size, device) + + if use_tvm: + + assert len(set(self.window_size)) == 1, "Only constant window size is supported." + q_k_mask = get_q_k(input_size, inner_size, window_size[0], device) + k_q_mask = get_k_q(q_k_mask) + self.layers = nn.ModuleList([ + EncoderLayer(d_model, d_inner_hid, num_heads, d_k, d_v, dropout=dropout, \ + normalize_before=False, use_tvm=True, q_k_mask=q_k_mask, k_q_mask=k_q_mask) for i in range(n_layer) + ]) + else: + self.layers = nn.ModuleList([ + EncoderLayer(d_model, d_inner_hid, num_heads, d_k, d_v, dropout=dropout, \ + normalize_before=False) for i in range(n_layer) + ]) + + self.embedding = SingleStepEmbedding(covariate_size, num_seq, d_model, input_size, device) + + self.conv_layers = Bottleneck_Construct(d_model, window_size, d_k) + + def forward(self, sequence): + + seq_enc = self.embedding(sequence) + mask = self.mask.repeat(len(seq_enc), self.num_heads, 1, 1).to(sequence.device) + + seq_enc = self.conv_layers(seq_enc) + + for i in range(len(self.layers)): + seq_enc, _ = self.layers[i](seq_enc, mask) + + indexes = self.indexes.repeat(seq_enc.size(0), 1, 1, seq_enc.size(2)).to(seq_enc.device) + indexes = indexes.view(seq_enc.size(0), -1, seq_enc.size(2)) + all_enc = torch.gather(seq_enc, 1, indexes) + all_enc = all_enc.view(seq_enc.size(0), self.all_size[0], -1) + + return all_enc + +class PyraformerSSModel(nn.Module): + @validated() + def __init__(self, freq, covariate_size, + num_seq, input_size ,dropout , d_model, + d_inner_hid, d_k, d_v , + num_heads , n_layer, loss , + window_size , inner_size, + use_tvm, prediction_length,context_length,lags_seq, + num_feat_dynamic_real, + num_feat_static_cat, + num_feat_static_real, + cardinality, + embedding_dimension, + distr_output, + # loss: DistributionLoss = NegativeLogLikelihood(), + scaling,num_parallel_samples,device): + + + super().__init__() + self.context_length = context_length + self.lags_seq = lags_seq or get_lags_for_frequency(freq_str=freq) + self.encoder = EncoderSS(covariate_size, + num_seq, input_size ,dropout , d_model, + d_inner_hid, d_k, d_v , + num_heads , n_layer, loss , + window_size , inner_size, + use_tvm, prediction_length, device) + + # convert hidden vectors into two scalar + self.mean_hidden = Predictor(4 * d_model, 1) + self.var_hidden = Predictor(4 * d_model, 1) + + self.softplus = nn.Softplus() + + def forward(self, data): + enc_output = self.encoder(data) + + mean_pre = self.mean_hidden(enc_output) + var_hid = self.var_hidden(enc_output) + var_pre = self.softplus(var_hid) + mean_pre = self.softplus(mean_pre) + + return mean_pre.squeeze(2), var_pre.squeeze(2) + + def test(self, data, v): + mu, sigma = self(data) + + sample_mu = mu[:, -1] * v + sample_sigma = sigma[:, -1] * v + return sample_mu, sample_sigma + + @property + def _past_length(self) -> int: + return self.context_length + max(self.lags_seq) + @property + def _number_of_features(self) -> int: + return ( + sum(self.embedding_dimension) + + self.num_feat_dynamic_real + + self.num_feat_static_real + + 1 # the log(scale) + ) + def get_lagged_subsequences( + self, sequence: torch.Tensor, subsequences_length: int, shift: int = 0 + ) -> 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). + subsequences_length : int + length of the subsequences to be extracted. + shift: int + shift the lags by this amount back. + 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, :]. + """ + sequence_length = sequence.shape[1] + indices = [lag - shift for lag in self.lags_seq] + + 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}" + ) + + 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 _check_shapes( + self, + prior_input: torch.Tensor, + inputs: torch.Tensor, + features: Optional[torch.Tensor], + ) -> None: + assert len(prior_input.shape) == len(inputs.shape) + assert ( + len(prior_input.shape) == 2 and self.input_size == 1 + ) or prior_input.shape[2] == self.input_size + assert (len(inputs.shape) == 2 and self.input_size == 1) or inputs.shape[ + -1 + ] == self.input_size + assert ( + features is None or features.shape[2] == self._number_of_features + ), f"{features.shape[2]}, expected {self._number_of_features}" + + def create_network_inputs( + 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: Optional[torch.Tensor] = None, + future_target: Optional[torch.Tensor] = None, + ): + # time feature + time_feat = ( + torch.cat( + ( + past_time_feat[:, self._past_length - self.context_length :, ...], + future_time_feat, + ), + dim=1, + ) + if future_target is not None + else past_time_feat[:, self._past_length - self.context_length :, ...] + ) + + # target + context = past_target[:, -self.context_length :] + observed_context = past_observed_values[:, -self.context_length :] + _, scale = self.scaler(context, observed_context) + + inputs = ( + torch.cat((past_target, future_target), dim=1) / scale + if future_target is not None + else past_target / scale + ) + + inputs_length = ( + self._past_length + self.prediction_length + if future_target is not None + else self._past_length + ) + assert inputs.shape[1] == inputs_length + + subsequences_length = ( + self.context_length + self.prediction_length + if future_target is not None + else self.context_length + ) + + # embeddings + embedded_cat = self.embedder(feat_static_cat) + static_feat = torch.cat( + (embedded_cat, feat_static_real, scale.log()), + dim=1, + ) + expanded_static_feat = static_feat.unsqueeze(1).expand( + -1, time_feat.shape[1], -1 + ) + + features = torch.cat((expanded_static_feat, time_feat), dim=-1) + + # self._check_shapes(prior_input, inputs, features) + + # sequence = torch.cat((prior_input, inputs), dim=1) + lagged_sequence = self.get_lagged_subsequences( + sequence=inputs, + subsequences_length=subsequences_length, + ) + + lags_shape = lagged_sequence.shape + reshaped_lagged_sequence = lagged_sequence.reshape( + lags_shape[0], lags_shape[1], -1 + ) + + transformer_inputs = torch.cat((reshaped_lagged_sequence, features), dim=-1) + + return transformer_inputs, scale, static_feat + + def output_params(self, transformer_inputs): + enc_input = transformer_inputs[:, : self.context_length, ...] + dec_input = transformer_inputs[:, self.context_length :, ...] + + enc_out = self.transformer.encoder(enc_input) + dec_output = self.transformer.decoder( + dec_input, enc_out, tgt_mask=self.tgt_mask + ) + + return self.param_proj(dec_output) + + @torch.jit.ignore + def output_distribution( + self, params, scale=None, trailing_n=None + ) -> torch.distributions.Distribution: + sliced_params = params + if trailing_n is not None: + sliced_params = [p[:, -trailing_n:] for p in params] + return self.distr_output.distribution(sliced_params, scale=scale) +class Encoder(nn.Module): + """ A encoder model with self attention mechanism. """ + + def __init__(self, model, window_size, truncate, input_size, inner_size,decoder,num_head,d_model, d_k,d_v,d_inner_hid, dropout,n_layer, ,enc_in,covariate_size,seq_num,CSCM,d_bottleneck,num_head,use_tvm,device): + super().__init__() + + self.d_model = d_model + self.model_type = model + self.window_size = window_size + self.truncate = truncate + if decoder == 'attention': + self.mask, self.all_size = get_mask(input_size, window_size, inner_size, device) + else: + self.mask, self.all_size = get_mask(input_size+1, window_size, inner_size, device) + self.decoder_type = decoder + if decoder == 'FC': + self.indexes = refer_points(self.all_size, window_size, device) + + if use_tvm: + assert len(set(self.window_size)) == 1, "Only constant window size is supported." + padding = 1 if decoder == 'FC' else 0 + q_k_mask = get_q_k(input_size + padding, inner_size, window_size[0],device) + k_q_mask = get_k_q(q_k_mask) + self.layers = nn.ModuleList([ + EncoderLayer(d_model, d_inner_hid, num_head, d_k, d_v, dropout=dropout, \ + normalize_before=False, use_tvm=True, q_k_mask=q_k_mask, k_q_mask=k_q_mask) for i in range(n_layer) + ]) + else: + self.layers = nn.ModuleList([ + EncoderLayer(d_model, d_inner_hid, num_head, d_k, d_v, dropout=dropout, \ + normalize_before=False) for i in range(n_layer) + ]) + + if opt.embed_type == 'CustomEmbedding': + self.enc_embedding = CustomEmbedding(enc_in, d_model, covariate_size, seq_num, dropout) + else: + self.enc_embedding = DataEmbedding(enc_in, d_model, dropout) + + self.conv_layers = eval(CSCM)(d_model, window_size, d_bottleneck) + + + def forward(self, x_enc, x_mark_enc): + seq_enc = self.enc_embedding(x_enc, x_mark_enc) + + mask = self.mask.repeat(len(seq_enc), 1, 1).to(x_enc.device) + seq_enc = self.conv_layers(seq_enc) + + for i in range(len(self.layers)): + seq_enc, _ = self.layers[i](seq_enc, mask) + + if self.decoder_type == 'FC': + indexes = self.indexes.repeat(seq_enc.size(0), 1, 1, seq_enc.size(2)).to(seq_enc.device) + indexes = indexes.view(seq_enc.size(0), -1, seq_enc.size(2)) + all_enc = torch.gather(seq_enc, 1, indexes) + seq_enc = all_enc.view(seq_enc.size(0), self.all_size[0], -1) + elif self.decoder_type == 'attention' and self.truncate: + seq_enc = seq_enc[:, :self.all_size[0]] + + return seq_enc + + +class PyraformerLRModel(nn.Module): + @validated() + def __init__(self, predict_step, d_model, input_size, decoder, window_size, truncate, model,d_inner_hid,num_head,d_k,d_v,dropout,enc_in,covariate_size,seq_num,CSCM,d_bottleneck,num_head,use_tvm,device): + super().__init__() + + self.predict_step = predict_step + self.d_model = d_model + self.input_size = input_size + self.decoder_type = decoder + self.channels = enc_in + + self.encoder = Encoder(model, window_size, truncate, input_size, inner_size,decoder,d_model, d_k,d_v,d_inner_hid, dropout,n_layer, ,enc_in,covariate_size,seq_num,CSCM,d_bottleneck,num_head,use_tvm,device) + if decoder == 'attention': + mask = get_subsequent_mask(input_size, window_size, predict_step, truncate) + self.decoder = Decoder(model,d_model,d_inner_hid,num_head,d_k,d_v,dropout,enc_in,covariate_size,seq_num, mask) + self.predictor = Predictor(d_model, enc_in) + elif opt.decoder == 'FC': + self.predictor = Predictor(4 * d_model, predict_step * enc_in) + + def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, pretrain): + """ + Return the hidden representations and predictions. + For a sequence (l_1, l_2, ..., l_N), we predict (l_2, ..., l_N, l_{N+1}). + Input: event_type: batch*seq_len; + event_time: batch*seq_len. + Output: enc_output: batch*seq_len*model_dim; + type_prediction: batch*seq_len*num_classes (not normalized); + time_prediction: batch*seq_len. + """ + if self.decoder_type == 'attention': + enc_output = self.encoder(x_enc, x_mark_enc) + dec_enc = self.decoder(x_dec, x_mark_dec, enc_output) + + if pretrain: + dec_enc = torch.cat([enc_output[:, :self.input_size], dec_enc], dim=1) + pred = self.predictor(dec_enc) + else: + pred = self.predictor(dec_enc) + elif self.decoder_type == 'FC': + enc_output = self.encoder(x_enc, x_mark_enc)[:, -1, :] + pred = self.predictor(enc_output).view(enc_output.size(0), self.predict_step, -1) + + return pred + +class TransformerModel(nn.Module): + @validated() + def __init__( + self, + freq: str, + context_length: int, + prediction_length: int, + num_feat_dynamic_real: int, + num_feat_static_real: int, + num_feat_static_cat: int, + cardinality: List[int], + # transformer arguments + nhead: int, + num_encoder_layers: int, + num_decoder_layers: int, + dim_feedforward: int, + activation: str = "gelu", + dropout: float = 0.1, + # univariate input + input_size: int = 1, + embedding_dimension: Optional[List[int]] = None, + distr_output: DistributionOutput = StudentTOutput(), + lags_seq: Optional[List[int]] = None, + scaling: bool = True, + num_parallel_samples: int = 100, + ) -> None: + super().__init__() + + self.input_size = input_size + + self.target_shape = distr_output.event_shape + self.num_feat_dynamic_real = num_feat_dynamic_real + self.num_feat_static_cat = num_feat_static_cat + self.num_feat_static_real = num_feat_static_real + self.embedding_dimension = ( + embedding_dimension + if embedding_dimension is not None or cardinality is None + else [min(50, (cat + 1) // 2) for cat in cardinality] + ) + self.lags_seq = lags_seq or get_lags_for_frequency(freq_str=freq) + self.num_parallel_samples = num_parallel_samples + self.history_length = context_length + max(self.lags_seq) + self.embedder = FeatureEmbedder( + cardinalities=cardinality, + embedding_dims=self.embedding_dimension, + ) + if scaling: + self.scaler = MeanScaler(dim=1, keepdim=True) + else: + self.scaler = NOPScaler(dim=1, keepdim=True) + + # total feature size + d_model = self.input_size * len(self.lags_seq) + self._number_of_features + + self.context_length = context_length + self.prediction_length = prediction_length + self.distr_output = distr_output + self.param_proj = distr_output.get_args_proj(d_model) + + # transformer enc-decoder and mask initializer + self.transformer = nn.Transformer( + d_model=d_model, + nhead=nhead, + num_encoder_layers=num_encoder_layers, + num_decoder_layers=num_decoder_layers, + dim_feedforward=dim_feedforward, + dropout=dropout, + activation=activation, + batch_first=True, + ) + + # causal decoder tgt mask + self.register_buffer( + "tgt_mask", + self.transformer.generate_square_subsequent_mask(prediction_length), + ) + + @property + def _number_of_features(self) -> int: + return ( + sum(self.embedding_dimension) + + self.num_feat_dynamic_real + + self.num_feat_static_real + + 1 # the log(scale) + ) + + @property + def _past_length(self) -> int: + return self.context_length + max(self.lags_seq) + + def get_lagged_subsequences( + self, sequence: torch.Tensor, subsequences_length: int, shift: int = 0 + ) -> 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). + subsequences_length : int + length of the subsequences to be extracted. + shift: int + shift the lags by this amount back. + 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, :]. + """ + sequence_length = sequence.shape[1] + indices = [lag - shift for lag in self.lags_seq] + + 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}" + ) + + 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 _check_shapes( + self, + prior_input: torch.Tensor, + inputs: torch.Tensor, + features: Optional[torch.Tensor], + ) -> None: + assert len(prior_input.shape) == len(inputs.shape) + assert ( + len(prior_input.shape) == 2 and self.input_size == 1 + ) or prior_input.shape[2] == self.input_size + assert (len(inputs.shape) == 2 and self.input_size == 1) or inputs.shape[ + -1 + ] == self.input_size + assert ( + features is None or features.shape[2] == self._number_of_features + ), f"{features.shape[2]}, expected {self._number_of_features}" + + def create_network_inputs( + 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: Optional[torch.Tensor] = None, + future_target: Optional[torch.Tensor] = None, + ): + # time feature + time_feat = ( + torch.cat( + ( + past_time_feat[:, self._past_length - self.context_length :, ...], + future_time_feat, + ), + dim=1, + ) + if future_target is not None + else past_time_feat[:, self._past_length - self.context_length :, ...] + ) + + # target + context = past_target[:, -self.context_length :] + observed_context = past_observed_values[:, -self.context_length :] + _, scale = self.scaler(context, observed_context) + + inputs = ( + torch.cat((past_target, future_target), dim=1) / scale + if future_target is not None + else past_target / scale + ) + + inputs_length = ( + self._past_length + self.prediction_length + if future_target is not None + else self._past_length + ) + assert inputs.shape[1] == inputs_length + + subsequences_length = ( + self.context_length + self.prediction_length + if future_target is not None + else self.context_length + ) + + # embeddings + embedded_cat = self.embedder(feat_static_cat) + static_feat = torch.cat( + (embedded_cat, feat_static_real, scale.log()), + dim=1, + ) + expanded_static_feat = static_feat.unsqueeze(1).expand( + -1, time_feat.shape[1], -1 + ) + + features = torch.cat((expanded_static_feat, time_feat), dim=-1) + + # self._check_shapes(prior_input, inputs, features) + + # sequence = torch.cat((prior_input, inputs), dim=1) + lagged_sequence = self.get_lagged_subsequences( + sequence=inputs, + subsequences_length=subsequences_length, + ) + + lags_shape = lagged_sequence.shape + reshaped_lagged_sequence = lagged_sequence.reshape( + lags_shape[0], lags_shape[1], -1 + ) + + transformer_inputs = torch.cat((reshaped_lagged_sequence, features), dim=-1) + + return transformer_inputs, scale, static_feat + + def output_params(self, transformer_inputs): + enc_input = transformer_inputs[:, : self.context_length, ...] + dec_input = transformer_inputs[:, self.context_length :, ...] + + enc_out = self.transformer.encoder(enc_input) + dec_output = self.transformer.decoder( + dec_input, enc_out, tgt_mask=self.tgt_mask + ) + + return self.param_proj(dec_output) + + @torch.jit.ignore + def output_distribution( + self, params, scale=None, trailing_n=None + ) -> torch.distributions.Distribution: + sliced_params = params + if trailing_n is not None: + sliced_params = [p[:, -trailing_n:] for p in params] + return self.distr_output.distribution(sliced_params, scale=scale) + + # for prediction + 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, + num_parallel_samples: Optional[int] = None, + ) -> torch.Tensor: + + if num_parallel_samples is None: + num_parallel_samples = self.num_parallel_samples + + encoder_inputs, scale, static_feat = self.create_network_inputs( + feat_static_cat, + feat_static_real, + past_time_feat, + past_target, + past_observed_values, + ) + + enc_out = self.transformer.encoder(encoder_inputs) + + repeated_scale = scale.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + + repeated_past_target = ( + past_target.repeat_interleave(repeats=self.num_parallel_samples, dim=0) + / repeated_scale + ) + + expanded_static_feat = static_feat.unsqueeze(1).expand( + -1, future_time_feat.shape[1], -1 + ) + features = torch.cat((expanded_static_feat, future_time_feat), dim=-1) + repeated_features = features.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + + repeated_enc_out = enc_out.repeat_interleave( + repeats=self.num_parallel_samples, dim=0 + ) + + future_samples = [] + + # greedy decoding + for k in range(self.prediction_length): + # self._check_shapes(repeated_past_target, next_sample, next_features) + # sequence = torch.cat((repeated_past_target, next_sample), dim=1) + + lagged_sequence = self.get_lagged_subsequences( + sequence=repeated_past_target, + subsequences_length=1 + k, + shift=1, + ) + + lags_shape = lagged_sequence.shape + reshaped_lagged_sequence = lagged_sequence.reshape( + lags_shape[0], lags_shape[1], -1 + ) + + decoder_input = torch.cat( + (reshaped_lagged_sequence, repeated_features[:, : k + 1]), dim=-1 + ) + + output = self.transformer.decoder(decoder_input, repeated_enc_out) + + params = self.param_proj(output[:, -1:]) + distr = self.output_distribution(params, scale=repeated_scale) + next_sample = distr.sample() + + repeated_past_target = torch.cat( + (repeated_past_target, next_sample / repeated_scale), dim=1 + ) + future_samples.append(next_sample) + + concat_future_samples = torch.cat(future_samples, dim=1) + return concat_future_samples.reshape( + (-1, self.num_parallel_samples, self.prediction_length) + self.target_shape, + ) diff --git a/Pyraformer/pyraformer.ipynb b/Pyraformer/pyraformer.ipynb new file mode 100644 index 0000000..77034aa --- /dev/null +++ b/Pyraformer/pyraformer.ipynb @@ -0,0 +1 @@ +{"cells":[{"cell_type":"code","execution_count":1,"id":"b19f0e22","metadata":{"id":"b19f0e22","executionInfo":{"status":"ok","timestamp":1649741813730,"user_tz":240,"elapsed":4,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}}},"outputs":[],"source":["%matplotlib inline"]},{"cell_type":"code","source":["# from google.colab import drive\n","# drive.mount('/content/drive/')\n","#%cd /content/drive/MyDrive/Udem/Sem2/Representation_Learning/IFT6135_Programming/Pyraformer/transformer"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"M6fjjCc2w6rX","executionInfo":{"status":"ok","timestamp":1649741815324,"user_tz":240,"elapsed":1003,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}},"outputId":"cb28e0cc-c6e0-45e3-ba60-79566978e080"},"id":"M6fjjCc2w6rX","execution_count":2,"outputs":[{"output_type":"stream","name":"stdout","text":["Drive already mounted at /content/drive/; to attempt to forcibly remount, call drive.mount(\"/content/drive/\", force_remount=True).\n"]}]},{"cell_type":"code","source":["!pip install pytorch-lightning"],"metadata":{"id":"_a4EOr95gtxR","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1649741259354,"user_tz":240,"elapsed":8842,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}},"outputId":"5c7e860c-6d26-4248-d2f3-7f58804506c9"},"id":"_a4EOr95gtxR","execution_count":4,"outputs":[{"output_type":"stream","name":"stdout","text":["Collecting pytorch-lightning\n"," Downloading pytorch_lightning-1.6.0-py3-none-any.whl (582 kB)\n","\u001b[?25l\r\u001b[K |▋ | 10 kB 17.4 MB/s eta 0:00:01\r\u001b[K |█▏ | 20 kB 21.1 MB/s eta 0:00:01\r\u001b[K |█▊ | 30 kB 23.0 MB/s eta 0:00:01\r\u001b[K |██▎ | 40 kB 15.8 MB/s eta 0:00:01\r\u001b[K |██▉ | 51 kB 11.9 MB/s eta 0:00:01\r\u001b[K |███▍ | 61 kB 13.5 MB/s eta 0:00:01\r\u001b[K |████ | 71 kB 14.7 MB/s eta 0:00:01\r\u001b[K |████▌ | 81 kB 12.8 MB/s eta 0:00:01\r\u001b[K |█████ | 92 kB 13.9 MB/s eta 0:00:01\r\u001b[K |█████▋ | 102 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████▏ | 112 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████▊ | 122 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████▎ | 133 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████▉ | 143 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████▍ | 153 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████ | 163 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████▋ | 174 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████▏ | 184 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████▊ | 194 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████▎ | 204 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████▉ | 215 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████▍ | 225 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████ | 235 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████▌ | 245 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████ | 256 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████▋ | 266 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████▏ | 276 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████▊ | 286 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████▎ | 296 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████▉ | 307 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████████▌ | 317 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████ | 327 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████▋ | 337 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████▏ | 348 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████▊ | 358 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████▎ | 368 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████▉ | 378 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████████████▍ | 389 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████████ | 399 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████████▌ | 409 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████ | 419 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████▋ | 430 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████████▏ | 440 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████████▊ | 450 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████████████████▎ | 460 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████████████ | 471 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████████████▌ | 481 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████████ | 491 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████████▋ | 501 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████████████▏ | 512 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████████████▊ | 522 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████████████████████▎ | 532 kB 14.3 MB/s eta 0:00:01\r\u001b[K |█████████████████████████████▉ | 542 kB 14.3 MB/s eta 0:00:01\r\u001b[K |██████████████████████████████▍ | 552 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████████████ | 563 kB 14.3 MB/s eta 0:00:01\r\u001b[K |███████████████████████████████▌| 573 kB 14.3 MB/s eta 0:00:01\r\u001b[K |████████████████████████████████| 582 kB 14.3 MB/s \n","\u001b[?25hRequirement already satisfied: torch>=1.8.* in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (1.10.0+cu111)\n","Requirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (4.1.1)\n","Requirement already satisfied: fsspec[http]!=2021.06.0,>=2021.05.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (2022.3.0)\n","Requirement already satisfied: numpy>=1.17.2 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (1.21.5)\n","Requirement already satisfied: tensorboard>=2.2.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (2.8.0)\n","Collecting torchmetrics>=0.4.1\n"," Downloading torchmetrics-0.7.3-py3-none-any.whl (398 kB)\n","\u001b[K |████████████████████████████████| 398 kB 38.5 MB/s \n","\u001b[?25hCollecting PyYAML>=5.4\n"," Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n","\u001b[K |████████████████████████████████| 596 kB 40.5 MB/s \n","\u001b[?25hRequirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (21.3)\n","Collecting pyDeprecate<0.4.0,>=0.3.1\n"," Downloading pyDeprecate-0.3.2-py3-none-any.whl (10 kB)\n","Requirement already satisfied: tqdm>=4.41.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (4.63.0)\n","Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.23.0)\n","Requirement already satisfied: aiohttp in /usr/local/lib/python3.7/dist-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.8.1)\n","Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=17.0->pytorch-lightning) (3.0.7)\n","Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (57.4.0)\n","Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.35.0)\n","Requirement already satisfied: protobuf>=3.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.17.3)\n","Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.44.0)\n","Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.0.1)\n","Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.8.1)\n","Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.6.1)\n","Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.4.6)\n","Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.3.6)\n","Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.37.1)\n","Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.0.0)\n","Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.4->tensorboard>=2.2.0->pytorch-lightning) (1.15.0)\n","Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.2.8)\n","Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.8)\n","Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.2.4)\n","Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (1.3.1)\n","Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard>=2.2.0->pytorch-lightning) (4.11.3)\n","Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard>=2.2.0->pytorch-lightning) (3.7.0)\n","Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.4.8)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.0.4)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.25.11)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.10)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2021.10.8)\n","Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (3.2.0)\n","Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (21.4.0)\n","Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (4.0.2)\n","Requirement already satisfied: asynctest==0.13.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (0.13.0)\n","Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.0.12)\n","Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.7.2)\n","Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.3.0)\n","Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (6.0.2)\n","Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.2.0)\n","Installing collected packages: pyDeprecate, torchmetrics, PyYAML, pytorch-lightning\n"," Attempting uninstall: PyYAML\n"," Found existing installation: PyYAML 3.13\n"," Uninstalling PyYAML-3.13:\n"," Successfully uninstalled PyYAML-3.13\n","Successfully installed PyYAML-6.0 pyDeprecate-0.3.2 pytorch-lightning-1.6.0 torchmetrics-0.7.3\n"]}]},{"cell_type":"code","source":["!pip install gluonts\n","!pip install datasets"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"k8r9pwTMhpuZ","executionInfo":{"status":"ok","timestamp":1649741250517,"user_tz":240,"elapsed":20780,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}},"outputId":"dd03cba5-e1fc-4472-8602-02dac453d8b1"},"id":"k8r9pwTMhpuZ","execution_count":3,"outputs":[{"output_type":"stream","name":"stdout","text":["Collecting gluonts\n"," Downloading gluonts-0.9.2-py3-none-any.whl (2.8 MB)\n","\u001b[K |████████████████████████████████| 2.8 MB 28.1 MB/s \n","\u001b[?25hCollecting pydantic~=1.1\n"," Downloading pydantic-1.9.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (10.9 MB)\n","\u001b[K |████████████████████████████████| 10.9 MB 49.3 MB/s \n","\u001b[?25hRequirement already satisfied: toolz~=0.10 in /usr/local/lib/python3.7/dist-packages (from gluonts) (0.11.2)\n","Requirement already satisfied: numpy~=1.16 in /usr/local/lib/python3.7/dist-packages (from gluonts) (1.21.5)\n","Requirement already satisfied: pandas~=1.0 in /usr/local/lib/python3.7/dist-packages (from gluonts) (1.3.5)\n","Requirement already satisfied: holidays>=0.9 in /usr/local/lib/python3.7/dist-packages (from gluonts) (0.10.5.2)\n","Requirement already satisfied: tqdm~=4.23 in /usr/local/lib/python3.7/dist-packages (from gluonts) (4.63.0)\n","Requirement already satisfied: matplotlib~=3.0 in /usr/local/lib/python3.7/dist-packages (from gluonts) (3.2.2)\n","Collecting typing-extensions~=4.0\n"," Downloading typing_extensions-4.1.1-py3-none-any.whl (26 kB)\n","Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from holidays>=0.9->gluonts) (1.15.0)\n","Requirement already satisfied: python-dateutil in /usr/local/lib/python3.7/dist-packages (from holidays>=0.9->gluonts) (2.8.2)\n","Requirement already satisfied: hijri-converter in /usr/local/lib/python3.7/dist-packages (from holidays>=0.9->gluonts) (2.2.3)\n","Requirement already satisfied: convertdate>=2.3.0 in /usr/local/lib/python3.7/dist-packages (from holidays>=0.9->gluonts) (2.4.0)\n","Requirement already satisfied: korean-lunar-calendar in /usr/local/lib/python3.7/dist-packages (from holidays>=0.9->gluonts) (0.2.1)\n","Requirement already satisfied: pymeeus<=1,>=0.3.13 in /usr/local/lib/python3.7/dist-packages (from convertdate>=2.3.0->holidays>=0.9->gluonts) (0.5.11)\n","Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib~=3.0->gluonts) (1.4.0)\n","Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib~=3.0->gluonts) (0.11.0)\n","Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib~=3.0->gluonts) (3.0.7)\n","Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas~=1.0->gluonts) (2018.9)\n","Installing collected packages: typing-extensions, pydantic, gluonts\n"," Attempting uninstall: typing-extensions\n"," Found existing installation: typing-extensions 3.10.0.2\n"," Uninstalling typing-extensions-3.10.0.2:\n"," Successfully uninstalled typing-extensions-3.10.0.2\n","\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n","tensorflow 2.8.0 requires tf-estimator-nightly==2.8.0.dev2021122109, which is not installed.\n","arviz 0.11.4 requires typing-extensions<4,>=3.7.4.3, but you have typing-extensions 4.1.1 which is incompatible.\u001b[0m\n","Successfully installed gluonts-0.9.2 pydantic-1.9.0 typing-extensions-4.1.1\n","Collecting datasets\n"," Downloading datasets-2.0.0-py3-none-any.whl (325 kB)\n","\u001b[K |████████████████████████████████| 325 kB 31.0 MB/s \n","\u001b[?25hRequirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.7/dist-packages (from datasets) (4.63.0)\n","Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from datasets) (21.3)\n","Collecting responses<0.19\n"," Downloading responses-0.18.0-py3-none-any.whl (38 kB)\n","Requirement already satisfied: dill in /usr/local/lib/python3.7/dist-packages (from datasets) (0.3.4)\n","Collecting aiohttp\n"," Downloading aiohttp-3.8.1-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n","\u001b[K |████████████████████████████████| 1.1 MB 48.2 MB/s \n","\u001b[?25hRequirement already satisfied: multiprocess in /usr/local/lib/python3.7/dist-packages (from datasets) (0.70.12.2)\n","Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from datasets) (4.11.3)\n","Collecting huggingface-hub<1.0.0,>=0.1.0\n"," Downloading huggingface_hub-0.5.1-py3-none-any.whl (77 kB)\n","\u001b[K |████████████████████████████████| 77 kB 6.8 MB/s \n","\u001b[?25hCollecting xxhash\n"," Downloading xxhash-3.0.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (212 kB)\n","\u001b[K |████████████████████████████████| 212 kB 61.1 MB/s \n","\u001b[?25hRequirement already satisfied: pyarrow>=5.0.0 in /usr/local/lib/python3.7/dist-packages (from datasets) (6.0.1)\n","Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from datasets) (1.21.5)\n","Collecting fsspec[http]>=2021.05.0\n"," Downloading fsspec-2022.3.0-py3-none-any.whl (136 kB)\n","\u001b[K |████████████████████████████████| 136 kB 59.6 MB/s \n","\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.7/dist-packages (from datasets) (2.23.0)\n","Requirement already satisfied: pandas in /usr/local/lib/python3.7/dist-packages (from datasets) (1.3.5)\n","Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from huggingface-hub<1.0.0,>=0.1.0->datasets) (3.6.0)\n","Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.7/dist-packages (from huggingface-hub<1.0.0,>=0.1.0->datasets) (4.1.1)\n","Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from huggingface-hub<1.0.0,>=0.1.0->datasets) (3.13)\n","Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->datasets) (3.0.7)\n","Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->datasets) (2021.10.8)\n","Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->datasets) (2.10)\n","Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->datasets) (3.0.4)\n","Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests>=2.19.0->datasets) (1.24.3)\n","Collecting urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1\n"," Downloading urllib3-1.25.11-py2.py3-none-any.whl (127 kB)\n","\u001b[K |████████████████████████████████| 127 kB 57.7 MB/s \n","\u001b[?25hCollecting multidict<7.0,>=4.5\n"," Downloading multidict-6.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (94 kB)\n","\u001b[K |████████████████████████████████| 94 kB 4.3 MB/s \n","\u001b[?25hRequirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->datasets) (21.4.0)\n","Collecting frozenlist>=1.1.1\n"," Downloading frozenlist-1.3.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (144 kB)\n","\u001b[K |████████████████████████████████| 144 kB 61.6 MB/s \n","\u001b[?25hRequirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->datasets) (2.0.12)\n","Collecting aiosignal>=1.1.2\n"," Downloading aiosignal-1.2.0-py3-none-any.whl (8.2 kB)\n","Collecting yarl<2.0,>=1.0\n"," Downloading yarl-1.7.2-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (271 kB)\n","\u001b[K |████████████████████████████████| 271 kB 50.8 MB/s \n","\u001b[?25hCollecting async-timeout<5.0,>=4.0.0a3\n"," Downloading async_timeout-4.0.2-py3-none-any.whl (5.8 kB)\n","Collecting asynctest==0.13.0\n"," Downloading asynctest-0.13.0-py3-none-any.whl (26 kB)\n","Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->datasets) (3.7.0)\n","Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas->datasets) (2.8.2)\n","Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas->datasets) (2018.9)\n","Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas->datasets) (1.15.0)\n","Installing collected packages: multidict, frozenlist, yarl, urllib3, asynctest, async-timeout, aiosignal, fsspec, aiohttp, xxhash, responses, huggingface-hub, datasets\n"," Attempting uninstall: urllib3\n"," Found existing installation: urllib3 1.24.3\n"," Uninstalling urllib3-1.24.3:\n"," Successfully uninstalled urllib3-1.24.3\n","\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n","datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\u001b[0m\n","Successfully installed aiohttp-3.8.1 aiosignal-1.2.0 async-timeout-4.0.2 asynctest-0.13.0 datasets-2.0.0 frozenlist-1.3.0 fsspec-2022.3.0 huggingface-hub-0.5.1 multidict-6.0.2 responses-0.18.0 urllib3-1.25.11 xxhash-3.0.0 yarl-1.7.2\n"]}]},{"cell_type":"code","source":["%matplotlib inline\n","from matplotlib import pyplot as plt\n","import matplotlib.dates as mdates\n","\n","from itertools import islice"],"metadata":{"id":"1XLYCBAswBhQ","executionInfo":{"status":"ok","timestamp":1649748115686,"user_tz":240,"elapsed":155,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}}},"id":"1XLYCBAswBhQ","execution_count":29,"outputs":[]},{"cell_type":"code","source":["from gluonts.evaluation import make_evaluation_predictions, Evaluator\n","from gluonts.dataset.repository.datasets import get_dataset\n","\n","from estimator import PyraformerEstimator"],"metadata":{"id":"n0nOWRF-wFl2","executionInfo":{"status":"ok","timestamp":1649748116654,"user_tz":240,"elapsed":231,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}}},"id":"n0nOWRF-wFl2","execution_count":30,"outputs":[]},{"cell_type":"code","source":["dataset = get_dataset(\"electricity\")"],"metadata":{"id":"Qzi9eE6q7x5y","executionInfo":{"status":"ok","timestamp":1649748117551,"user_tz":240,"elapsed":143,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}}},"id":"Qzi9eE6q7x5y","execution_count":31,"outputs":[]},{"cell_type":"code","source":["estimator = PyraformerEstimator(\n"," freq=dataset.metadata.freq,\n"," prediction_length=dataset.metadata.prediction_length,\n"," num_feat_static_cat=1,\n"," cardinality=[321],\n","\n"," batch_size=1,\n"," num_batches_per_epoch=100,\n"," trainer_kwargs=dict(max_epochs=1, accelerator='gpu', gpus=1),\n",")"],"metadata":{"id":"i7AV93A07sQa","executionInfo":{"status":"ok","timestamp":1649748118628,"user_tz":240,"elapsed":203,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}}},"id":"i7AV93A07sQa","execution_count":32,"outputs":[]},{"cell_type":"code","source":["predictor = estimator.train(training_data=dataset.train,num_workers=8)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":355},"id":"hXjDU6rhK9H_","executionInfo":{"status":"error","timestamp":1649748120714,"user_tz":240,"elapsed":611,"user":{"displayName":"Hena Ghonia","userId":"03246241722682988409"}},"outputId":"a767d166-b3a0-4d4c-e67c-8c87a766b1dc"},"id":"hXjDU6rhK9H_","execution_count":33,"outputs":[{"output_type":"stream","name":"stderr","text":["/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/parsing.py:245: UserWarning: Attribute 'model' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['model'])`.\n"," f\"Attribute {k!r} is an instance of `nn.Module` and is already saved during checkpointing.\"\n","/usr/local/lib/python3.7/dist-packages/pytorch_lightning/utilities/parsing.py:245: UserWarning: Attribute 'loss' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['loss'])`.\n"," f\"Attribute {k!r} is an instance of `nn.Module` and is already saved during checkpointing.\"\n"]},{"output_type":"error","ename":"ValidationError","evalue":"ignored","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mValidationError\u001b[0m Traceback (most recent call last)","\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mpredictor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mestimator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtraining_data\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnum_workers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m8\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gluonts/torch/model/estimator.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, training_data, validation_data, num_workers, shuffle_buffer_length, cache_data, ckpt_path, **kwargs)\u001b[0m\n\u001b[1;32m 195\u001b[0m \u001b[0mcache_data\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcache_data\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 196\u001b[0m \u001b[0mckpt_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mckpt_path\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 197\u001b[0;31m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 198\u001b[0m ).predictor\n","\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gluonts/torch/model/estimator.py\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(self, training_data, validation_data, num_workers, shuffle_buffer_length, cache_data, ckpt_path, **kwargs)\u001b[0m\n\u001b[1;32m 127\u001b[0m \u001b[0mtraining_network\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[0mnum_workers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_workers\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 129\u001b[0;31m \u001b[0mshuffle_buffer_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mshuffle_buffer_length\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 130\u001b[0m )\n\u001b[1;32m 131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/content/drive/MyDrive/Udem/Sem2/Representation_Learning/IFT6135_Programming/Pyraformer/transformer/estimator.py\u001b[0m in \u001b[0;36mcreate_training_data_loader\u001b[0;34m(self, data, module, shuffle_buffer_length, **kwargs)\u001b[0m\n\u001b[1;32m 260\u001b[0m ) -> Iterable:\n\u001b[1;32m 261\u001b[0m transformation = self._create_instance_splitter(\n\u001b[0;32m--> 262\u001b[0;31m \u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"training\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 263\u001b[0m ) + SelectFields(TRAINING_INPUT_NAMES)\n\u001b[1;32m 264\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/content/drive/MyDrive/Udem/Sem2/Representation_Learning/IFT6135_Programming/Pyraformer/transformer/estimator.py\u001b[0m in \u001b[0;36m_create_instance_splitter\u001b[0;34m(self, module, mode)\u001b[0m\n\u001b[1;32m 249\u001b[0m \u001b[0mFieldName\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mOBSERVED_VALUES\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 250\u001b[0m ],\n\u001b[0;32m--> 251\u001b[0;31m \u001b[0mdummy_value\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdistr_output\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalue_in_support\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 252\u001b[0m )\n\u001b[1;32m 253\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.7/dist-packages/gluonts/core/component.py\u001b[0m in \u001b[0;36minit_wrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 323\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m\"self\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 324\u001b[0m }\n\u001b[0;32m--> 325\u001b[0;31m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mPydanticModel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0mnmargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 326\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 327\u001b[0m \u001b[0;31m# merge nmargs, kwargs, and the model fields into a single dict\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.7/dist-packages/pydantic/main.cpython-37m-x86_64-linux-gnu.so\u001b[0m in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n","\u001b[0;31mValidationError\u001b[0m: 1 validation error for InstanceSplitterModel\ninstance_sampler\n value is not a valid dict (type=type_error.dict)"]}]},{"cell_type":"code","execution_count":null,"id":"d61f32ab","metadata":{"id":"d61f32ab"},"outputs":[],"source":["# plt.figure(figsize=(20, 15))\n","# date_formater = mdates.DateFormatter('%b, %d')\n","# plt.rcParams.update({'font.size': 15})\n","\n","# for idx, (forecast, ts) in islice(enumerate(zip(forecasts, tss)), 9):\n","# ax = plt.subplot(3, 3, idx+1)\n","\n","# plt.plot(ts[-4 * dataset.metadata.prediction_length:], label=\"target\", )\n","# forecast.plot( color='g')\n","# plt.xticks(rotation=60)\n","# ax.xaxis.set_major_formatter(date_formater)\n","\n","# plt.gcf().tight_layout()\n","# plt.legend()\n","# plt.show()"]},{"cell_type":"code","execution_count":null,"id":"d494463f","metadata":{"id":"d494463f"},"outputs":[],"source":["# def plot_prob_forecasts(ts_entry, forecast_entry):\n","# plot_length = 70\n","# prediction_intervals = (50.0, 90.0)\n","# legend = [\"observations\", \"median prediction\"] + [f\"{k}% prediction interval\" for k in prediction_intervals][::-1]\n","\n","# fig, ax = plt.subplots(1, 1, figsize=(10, 7))\n","# ts_entry[-plot_length:].plot(ax=ax) # plot the time series\n","# forecast_entry.plot(prediction_intervals=prediction_intervals, color='g')\n","# plt.grid(which=\"both\")\n","# plt.legend(legend, loc=\"best\")\n","# plt.show()"]},{"cell_type":"code","execution_count":null,"id":"5256fde1","metadata":{"id":"5256fde1"},"outputs":[],"source":["# index = 123\n","# plot_prob_forecasts(tss[index], forecasts[index])"]},{"cell_type":"code","execution_count":null,"id":"66a41556","metadata":{"id":"66a41556"},"outputs":[],"source":[""]}],"metadata":{"kernelspec":{"display_name":"Python 3 (ipykernel)","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.9.7"},"colab":{"name":"pyraformer.ipynb","provenance":[],"collapsed_sections":[]}},"nbformat":4,"nbformat_minor":5} \ No newline at end of file diff --git a/Pyraformer/pyraformer/Layers.py b/Pyraformer/pyraformer/Layers.py new file mode 100644 index 0000000..2cd96af --- /dev/null +++ b/Pyraformer/pyraformer/Layers.py @@ -0,0 +1,408 @@ +from torch.functional import align_tensors +import torch.nn as nn + +from torch.nn.modules.linear import Linear +from .SubLayers import MultiHeadAttention, PositionwiseFeedForward +import torch +from .embed import DataEmbedding, CustomEmbedding +import math + + + +def get_mask(input_size, window_size, inner_size, device): + """Get the attention mask of PAM-Naive""" + # Get the size of all layers + all_size = [] + all_size.append(input_size) + for i in range(len(window_size)): + layer_size = math.floor(all_size[i] / window_size[i]) + all_size.append(layer_size) + + seq_length = sum(all_size) + mask = torch.zeros(seq_length, seq_length, device=device) + + # get intra-scale mask + inner_window = inner_size // 2 + for layer_idx in range(len(all_size)): + start = sum(all_size[:layer_idx]) + for i in range(start, start + all_size[layer_idx]): + left_side = max(i - inner_window, start) + right_side = min(i + inner_window + 1, start + all_size[layer_idx]) + mask[i, left_side:right_side] = 1 + + # get inter-scale mask + for layer_idx in range(1, len(all_size)): + start = sum(all_size[:layer_idx]) + for i in range(start, start + all_size[layer_idx]): + left_side = (start - all_size[layer_idx - 1]) + (i - start) * window_size[layer_idx - 1] + if i == ( start + all_size[layer_idx] - 1): + right_side = start + else: + right_side = (start - all_size[layer_idx - 1]) + (i - start + 1) * window_size[layer_idx - 1] + mask[i, left_side:right_side] = 1 + mask[left_side:right_side, i] = 1 + + mask = (1 - mask).bool() + + return mask, all_size + + +def refer_points(all_sizes, window_size, device): + """Gather features from PAM's pyramid sequences""" + input_size = all_sizes[0] + indexes = torch.zeros(input_size, len(all_sizes), device=device) + + for i in range(input_size): + indexes[i][0] = i + former_index = i + for j in range(1, len(all_sizes)): + start = sum(all_sizes[:j]) + inner_layer_idx = former_index - (start - all_sizes[j - 1]) + former_index = start + min(inner_layer_idx // window_size[j - 1], all_sizes[j] - 1) + indexes[i][j] = former_index + + indexes = indexes.unsqueeze(0).unsqueeze(3) + + return indexes.long() + + +def get_subsequent_mask(input_size, window_size, predict_step, truncate): + """Get causal attention mask for decoder.""" + if truncate: + mask = torch.zeros(predict_step, input_size + predict_step) + for i in range(predict_step): + mask[i][:input_size+i+1] = 1 + mask = (1 - mask).bool().unsqueeze(0) + else: + all_size = [] + all_size.append(input_size) + for i in range(len(window_size)): + layer_size = math.floor(all_size[i] / window_size[i]) + all_size.append(layer_size) + all_size = sum(all_size) + mask = torch.zeros(predict_step, all_size + predict_step) + for i in range(predict_step): + mask[i][:all_size+i+1] = 1 + mask = (1 - mask).bool().unsqueeze(0) + + return mask + + +def get_q_k(input_size, window_size, stride, device): + """ + Get the index of the key that a given query needs to attend to. + """ + second_length = input_size // stride + second_last = input_size - (second_length - 1) * stride + third_start = input_size + second_length + third_length = second_length // stride + third_last = second_length - (third_length - 1) * stride + max_attn = max(second_last, third_last) + fourth_start = third_start + third_length + fourth_length = third_length // stride + full_length = fourth_start + fourth_length + fourth_last = third_length - (fourth_length - 1) * stride + max_attn = max(third_last, fourth_last) + + max_attn += window_size + 1 + mask = torch.zeros(full_length, max_attn, dtype=torch.int32, device=device) - 1 + + for i in range(input_size): + mask[i, 0:window_size] = i + torch.arange(window_size) - window_size // 2 + mask[i, mask[i] > input_size - 1] = -1 + + mask[i, -1] = i // stride + input_size + mask[i][mask[i] > third_start - 1] = third_start - 1 + for i in range(second_length): + mask[input_size+i, 0:window_size] = input_size + i + torch.arange(window_size) - window_size // 2 + mask[input_size+i, mask[input_size+i] < input_size] = -1 + mask[input_size+i, mask[input_size+i] > third_start - 1] = -1 + + if i < second_length - 1: + mask[input_size+i, window_size:(window_size+stride)] = torch.arange(stride) + i * stride + else: + mask[input_size+i, window_size:(window_size+second_last)] = torch.arange(second_last) + i * stride + + mask[input_size+i, -1] = i // stride + third_start + mask[input_size+i, mask[input_size+i] > fourth_start - 1] = fourth_start - 1 + for i in range(third_length): + mask[third_start+i, 0:window_size] = third_start + i + torch.arange(window_size) - window_size // 2 + mask[third_start+i, mask[third_start+i] < third_start] = -1 + mask[third_start+i, mask[third_start+i] > fourth_start - 1] = -1 + + if i < third_length - 1: + mask[third_start+i, window_size:(window_size+stride)] = input_size + torch.arange(stride) + i * stride + else: + mask[third_start+i, window_size:(window_size+third_last)] = input_size + torch.arange(third_last) + i * stride + + mask[third_start+i, -1] = i // stride + fourth_start + mask[third_start+i, mask[third_start+i] > full_length - 1] = full_length - 1 + for i in range(fourth_length): + mask[fourth_start+i, 0:window_size] = fourth_start + i + torch.arange(window_size) - window_size // 2 + mask[fourth_start+i, mask[fourth_start+i] < fourth_start] = -1 + mask[fourth_start+i, mask[fourth_start+i] > full_length - 1] = -1 + + if i < fourth_length - 1: + mask[fourth_start+i, window_size:(window_size+stride)] = third_start + torch.arange(stride) + i * stride + else: + mask[fourth_start+i, window_size:(window_size+fourth_last)] = third_start + torch.arange(fourth_last) + i * stride + + return mask + + +def get_k_q(q_k_mask): + """ + Get the index of the query that can attend to the given key. + """ + k_q_mask = q_k_mask.clone() + for i in range(len(q_k_mask)): + for j in range(len(q_k_mask[0])): + if q_k_mask[i, j] >= 0: + k_q_mask[i, j] = torch.where(q_k_mask[q_k_mask[i, j]] ==i )[0] + + return k_q_mask + + +class EncoderLayer(nn.Module): + """ Compose with two layers """ + + def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True, use_tvm=False, q_k_mask=None, k_q_mask=None): + super(EncoderLayer, self).__init__() + self.use_tvm = use_tvm + if use_tvm: + from .PAM_TVM import PyramidalAttention + self.slf_attn = PyramidalAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before, q_k_mask=q_k_mask, k_q_mask=k_q_mask) + else: + self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before) + + self.pos_ffn = PositionwiseFeedForward( + d_model, d_inner, dropout=dropout, normalize_before=normalize_before) + + def forward(self, enc_input, slf_attn_mask=None): + if self.use_tvm: + enc_output = self.slf_attn(enc_input) + enc_slf_attn = None + else: + enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask) + + enc_output = self.pos_ffn(enc_output) + + return enc_output, enc_slf_attn + + +class DecoderLayer(nn.Module): + """ Compose with two layers """ + + def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True): + super(DecoderLayer, self).__init__() + self.slf_attn = MultiHeadAttention( + n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before) + self.pos_ffn = PositionwiseFeedForward( + d_model, d_inner, dropout=dropout, normalize_before=normalize_before) + + def forward(self, Q, K, V, slf_attn_mask=None): + enc_output, enc_slf_attn = self.slf_attn( + Q, K, V, mask=slf_attn_mask) + + enc_output = self.pos_ffn(enc_output) + + return enc_output, enc_slf_attn + + +class ConvLayer(nn.Module): + def __init__(self, c_in, window_size): + super(ConvLayer, self).__init__() + self.downConv = nn.Conv1d(in_channels=c_in, + out_channels=c_in, + kernel_size=window_size, + stride=window_size) + self.norm = nn.BatchNorm1d(c_in) + self.activation = nn.ELU() + + def forward(self, x): + x = self.downConv(x) + x = self.norm(x) + x = self.activation(x) + return x + + +class Conv_Construct(nn.Module): + """Convolution CSCM""" + def __init__(self, d_model, window_size, d_inner): + super(Conv_Construct, self).__init__() + if not isinstance(window_size, list): + self.conv_layers = nn.ModuleList([ + ConvLayer(d_model, window_size), + ConvLayer(d_model, window_size), + ConvLayer(d_model, window_size) + ]) + else: + self.conv_layers = nn.ModuleList([ + ConvLayer(d_model, window_size[0]), + ConvLayer(d_model, window_size[1]), + ConvLayer(d_model, window_size[2]) + ]) + self.norm = nn.LayerNorm(d_model) + + def forward(self, enc_input): + all_inputs = [] + enc_input = enc_input.permute(0, 2, 1) + all_inputs.append(enc_input) + + for i in range(len(self.conv_layers)): + enc_input = self.conv_layers[i](enc_input) + all_inputs.append(enc_input) + + all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2) + all_inputs = self.norm(all_inputs) + + return all_inputs + + +class Bottleneck_Construct(nn.Module): + """Bottleneck convolution CSCM""" + def __init__(self, d_model, window_size, d_inner): + super(Bottleneck_Construct, self).__init__() + if not isinstance(window_size, list): + self.conv_layers = nn.ModuleList([ + ConvLayer(d_inner, window_size), + ConvLayer(d_inner, window_size), + ConvLayer(d_inner, window_size) + ]) + else: + self.conv_layers = [] + for i in range(len(window_size)): + self.conv_layers.append(ConvLayer(d_inner, window_size[i])) + self.conv_layers = nn.ModuleList(self.conv_layers) + self.up = Linear(d_inner, d_model) + self.down = Linear(d_model, d_inner) + self.norm = nn.LayerNorm(d_model) + + def forward(self, enc_input): + + temp_input = self.down(enc_input).permute(0, 2, 1) + all_inputs = [] + for i in range(len(self.conv_layers)): + temp_input = self.conv_layers[i](temp_input) + all_inputs.append(temp_input) + + all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2) + all_inputs = self.up(all_inputs) + all_inputs = torch.cat([enc_input, all_inputs], dim=1) + + all_inputs = self.norm(all_inputs) + + return all_inputs + + +class MaxPooling_Construct(nn.Module): + """Max pooling CSCM""" + def __init__(self, d_model, window_size, d_inner): + super(MaxPooling_Construct, self).__init__() + if not isinstance(window_size, list): + self.pooling_layers = nn.ModuleList([ + nn.MaxPool1d(kernel_size=window_size), + nn.MaxPool1d(kernel_size=window_size), + nn.MaxPool1d(kernel_size=window_size) + ]) + else: + self.pooling_layers = nn.ModuleList([ + nn.MaxPool1d(kernel_size=window_size[0]), + nn.MaxPool1d(kernel_size=window_size[1]), + nn.MaxPool1d(kernel_size=window_size[2]) + ]) + self.norm = nn.LayerNorm(d_model) + + def forward(self, enc_input): + all_inputs = [] + enc_input = enc_input.transpose(1, 2).contiguous() + all_inputs.append(enc_input) + + for layer in self.pooling_layers: + enc_input = layer(enc_input) + all_inputs.append(enc_input) + + all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2) + all_inputs = self.norm(all_inputs) + + return all_inputs + + +class AvgPooling_Construct(nn.Module): + """Average pooling CSCM""" + def __init__(self, d_model, window_size, d_inner): + super(AvgPooling_Construct, self).__init__() + if not isinstance(window_size, list): + self.pooling_layers = nn.ModuleList([ + nn.AvgPool1d(kernel_size=window_size), + nn.AvgPool1d(kernel_size=window_size), + nn.AvgPool1d(kernel_size=window_size) + ]) + else: + self.pooling_layers = nn.ModuleList([ + nn.AvgPool1d(kernel_size=window_size[0]), + nn.AvgPool1d(kernel_size=window_size[1]), + nn.AvgPool1d(kernel_size=window_size[2]) + ]) + self.norm = nn.LayerNorm(d_model) + + def forward(self, enc_input): + all_inputs = [] + enc_input = enc_input.transpose(1, 2).contiguous() + all_inputs.append(enc_input) + + for layer in self.pooling_layers: + enc_input = layer(enc_input) + all_inputs.append(enc_input) + + all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2) + all_inputs = self.norm(all_inputs) + + return all_inputs + + +class Predictor(nn.Module): + + def __init__(self, dim, num_types): + super().__init__() + + self.linear = nn.Linear(dim, num_types, bias=False) + nn.init.xavier_normal_(self.linear.weight) + + def forward(self, data): + out = self.linear(data) + out = out + return out + + +class Decoder(nn.Module): + """ A encoder model with self attention mechanism. """ + + def __init__(self, model,d_model,d_inner_hid,num_head,d_k,d_v,dropout,enc_in,covariate_size,seq_num, mask): + super().__init__() + + self.model_type = model + self.mask = mask + + self.layers = nn.ModuleList([ + DecoderLayer(d_model, d_inner_hid, num_head, d_k, d_v, dropout=dropout, \ + normalize_before=False), + DecoderLayer(d_model, d_inner_hid, num_head, d_k, d_v, dropout=dropout, \ + normalize_before=False) + ]) + + if opt.embed_type == 'CustomEmbedding': + self.dec_embedding = CustomEmbedding(enc_in, d_model, covariate_size, seq_num, dropout) + else: + self.dec_embedding = DataEmbedding(enc_in, d_model, dropout) + + def forward(self, x_dec, x_mark_dec, refer): + dec_enc = self.dec_embedding(x_dec, x_mark_dec) + + dec_enc, _ = self.layers[0](dec_enc, refer, refer) + refer_enc = torch.cat([refer, dec_enc], dim=1) + mask = self.mask.repeat(len(dec_enc), 1, 1).to(dec_enc.device) + dec_enc, _ = self.layers[1](dec_enc, refer_enc, refer_enc, slf_attn_mask=mask) + + return dec_enc + diff --git a/Pyraformer/pyraformer/Modules.py b/Pyraformer/pyraformer/Modules.py new file mode 100644 index 0000000..d9b14bd --- /dev/null +++ b/Pyraformer/pyraformer/Modules.py @@ -0,0 +1,25 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class ScaledDotProductAttention(nn.Module): + """ Scaled Dot-Product Attention """ + + def __init__(self, temperature, attn_dropout=0.2): + super().__init__() + + self.temperature = temperature + self.dropout = nn.Dropout(attn_dropout) + + def forward(self, q, k, v, mask=None): + attn = torch.matmul(q / self.temperature, k.transpose(2, 3)) + + if mask is not None: + attn = attn.masked_fill(mask, -1e9) + + attn = self.dropout(F.softmax(attn, dim=-1)) + output = torch.matmul(attn, v) + + return output, attn + diff --git a/Pyraformer/pyraformer/PAM_TVM.py b/Pyraformer/pyraformer/PAM_TVM.py new file mode 100644 index 0000000..706a1a5 --- /dev/null +++ b/Pyraformer/pyraformer/PAM_TVM.py @@ -0,0 +1,65 @@ +import torch.nn as nn +import torch.nn.functional as F +import math +from .hierarchical_mm_tvm import graph_mm as graph_mm_tvm + + +class PyramidalAttention(nn.Module): + def __init__(self, n_head, d_model, d_k, d_v, dropout, normalize_before, q_k_mask, k_q_mask): + super(PyramidalAttention, self).__init__() + self.normalize_before = normalize_before + self.n_head = n_head + self.d_k = d_k + + self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) + self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False) + self.w_vs = nn.Linear(d_model, n_head * d_k, bias=False) + nn.init.xavier_uniform_(self.w_qs.weight) + nn.init.xavier_uniform_(self.w_ks.weight) + nn.init.xavier_uniform_(self.w_vs.weight) + + self.fc = nn.Linear(d_k * n_head, d_model) + nn.init.xavier_uniform_(self.fc.weight) + + self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) + self.dropout_attn = nn.Dropout(dropout) + self.dropout_fc = nn.Dropout(dropout) + self.q_k_mask = q_k_mask + self.k_q_mask = k_q_mask + + def forward(self, hidden_states): + residual = hidden_states + + hidden_states = hidden_states + bsz, seq_len, _ = hidden_states.size() + + q = hidden_states + if self.normalize_before: + q = self.layer_norm(q) + + q = self.w_qs(q) + k = self.w_ks(hidden_states) + v = self.w_vs(hidden_states) + q /= math.sqrt(self.d_k) + + q = q.view(bsz, seq_len, self.n_head, self.d_k) + k = k.view(bsz, seq_len, self.n_head, self.d_k) + q = q.float().contiguous() + k = k.float().contiguous() + # attn_weights.size(): (batch_size, L, num_heads, 11) + attn_weights = graph_mm_tvm(q, k, self.q_k_mask, self.k_q_mask, False, 0) + attn_weights = self.dropout_attn(F.softmax(attn_weights, dim=-1)) + + v = v.view(bsz, seq_len, self.n_head, self.d_k) + v = v.float().contiguous() + # is_t1_diagonaled=True + attn = graph_mm_tvm(attn_weights, v, self.q_k_mask, self.k_q_mask, True, 0) + attn = attn.reshape(bsz, seq_len, self.n_head * self.d_k).contiguous() + context = self.dropout_fc(self.fc(attn)) + context += residual + + if not self.normalize_before: + context = self.layer_norm(context) + + return context + diff --git a/Pyraformer/pyraformer/SubLayers.py b/Pyraformer/pyraformer/SubLayers.py new file mode 100644 index 0000000..e252d5d --- /dev/null +++ b/Pyraformer/pyraformer/SubLayers.py @@ -0,0 +1,96 @@ +import torch.nn as nn +import torch.nn.functional as F + +from .Modules import ScaledDotProductAttention + + +class MultiHeadAttention(nn.Module): + """ Multi-Head Attention module """ + + def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1, normalize_before=True): + super().__init__() + + self.normalize_before = normalize_before + self.n_head = n_head + self.d_k = d_k + self.d_v = d_v + + self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) + self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False) + self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False) + nn.init.xavier_uniform_(self.w_qs.weight) + nn.init.xavier_uniform_(self.w_ks.weight) + nn.init.xavier_uniform_(self.w_vs.weight) + + self.fc = nn.Linear(d_v * n_head, d_model) + nn.init.xavier_uniform_(self.fc.weight) + + self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5, attn_dropout=dropout) + + self.layer_norm = nn.LayerNorm(d_model, eps=1e-6) + self.dropout = nn.Dropout(dropout) + + def forward(self, q, k, v, mask=None): + d_k, d_v, n_head = self.d_k, self.d_v, self.n_head + sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1) + + residual = q + if self.normalize_before: + q = self.layer_norm(q) + + # Pass through the pre-attention projection: b x lq x (n*dv) + # Separate different heads: b x lq x n x dv + q = self.w_qs(q).view(sz_b, len_q, n_head, d_k) + k = self.w_ks(k).view(sz_b, len_k, n_head, d_k) + v = self.w_vs(v).view(sz_b, len_v, n_head, d_v) + + # Transpose for attention dot product: b x n x lq x dv + q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + + if mask is not None: + if len(mask.size()) == 3: + mask = mask.unsqueeze(1) # For head axis broadcasting. + + output, attn = self.attention(q, k, v, mask=mask) + + # Transpose to move the head dimension back: b x lq x n x dv + # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv) + output = output.transpose(1, 2).contiguous().view(sz_b, len_q, -1) + output = self.dropout(self.fc(output)) + output += residual + + if not self.normalize_before: + output = self.layer_norm(output) + return output, attn + + +class PositionwiseFeedForward(nn.Module): + """ Two-layer position-wise feed-forward neural network. """ + + def __init__(self, d_in, d_hid, dropout=0.1, normalize_before=True): + super().__init__() + + self.normalize_before = normalize_before + + self.w_1 = nn.Linear(d_in, d_hid) + self.w_2 = nn.Linear(d_hid, d_in) + + self.layer_norm = nn.LayerNorm(d_in, eps=1e-6) + #self.layer_norm = GraphNorm(d_in) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + residual = x + if self.normalize_before: + x = self.layer_norm(x) + + x = F.gelu(self.w_1(x)) + x = self.dropout(x) + x = self.w_2(x) + x = self.dropout(x) + x = x + residual + + if not self.normalize_before: + x = self.layer_norm(x) + return x + diff --git a/Pyraformer/pyraformer/embed.py b/Pyraformer/pyraformer/embed.py new file mode 100644 index 0000000..3ac8c08 --- /dev/null +++ b/Pyraformer/pyraformer/embed.py @@ -0,0 +1,158 @@ +""" +Modified based on Informer. +@inproceedings{haoyietal-informer-2021, + author = {Haoyi Zhou and Shanghang Zhang and Jieqi Peng and Shuai Zhang and Jianxin Li and + Hui Xiong and Wancai Zhang}, + title = {Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting}, + booktitle = {The Thirty-Fifth {AAAI} Conference on Artificial Intelligence, {AAAI} 2021, Virtual Conference}, + volume = {35}, number = {12}, pages = {11106--11115}, publisher = {{AAAI} Press}, year = {2021}, +} +""" + +import torch +import torch.nn as nn + +import math + + +class PositionalEmbedding(nn.Module): + def __init__(self, d_model, max_len=5000): + super(PositionalEmbedding, self).__init__() + # Compute the positional encodings once in log space. + pe = torch.zeros(max_len, d_model).float() + pe.require_grad = False + + position = torch.arange(0, max_len).float().unsqueeze(1) + div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp() + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + + pe = pe.unsqueeze(0) + self.register_buffer('pe', pe) + + def forward(self, x): + return self.pe[:, :x.size(1)] + +class TokenEmbedding(nn.Module): + def __init__(self, c_in, d_model): + super(TokenEmbedding, self).__init__() + padding = 1 if torch.__version__>='1.5.0' else 2 + self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model, + kernel_size=3, padding=padding, padding_mode='circular') + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight,mode='fan_in',nonlinearity='leaky_relu') + + def forward(self, x): + x = self.tokenConv(x.permute(0, 2, 1)).transpose(1,2) + return x + +class FixedEmbedding(nn.Module): + def __init__(self, c_in, d_model): + super(FixedEmbedding, self).__init__() + + w = torch.zeros(c_in, d_model).float() + w.require_grad = False + + position = torch.arange(0, c_in).float().unsqueeze(1) + div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp() + + w[:, 0::2] = torch.sin(position * div_term) + w[:, 1::2] = torch.cos(position * div_term) + + self.emb = nn.Embedding(c_in, d_model) + self.emb.weight = nn.Parameter(w, requires_grad=False) + + def forward(self, x): + return self.emb(x).detach() + +class TimeFeatureEmbedding(nn.Module): + def __init__(self, d_model): + super(TimeFeatureEmbedding, self).__init__() + + d_inp = 4 + self.embed = nn.Linear(d_inp, d_model) + + def forward(self, x): + return self.embed(x) + +"""Embedding modules. The DataEmbedding is used by the ETT dataset for long range forecasting.""" +class DataEmbedding(nn.Module): + def __init__(self, c_in, d_model, dropout=0.1): + super(DataEmbedding, self).__init__() + + self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model) + self.position_embedding = PositionalEmbedding(d_model=d_model) + self.temporal_embedding = TimeFeatureEmbedding(d_model) + + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + x = self.value_embedding(x) + self.position_embedding(x) + self.temporal_embedding(x_mark) + + return self.dropout(x) + +"""The CustomEmbedding is used by the electricity dataset and app flow dataset for long range forecasting.""" +class CustomEmbedding(nn.Module): + def __init__(self, c_in, d_model, temporal_size, seq_num, dropout=0.1): + super(CustomEmbedding, self).__init__() + + self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model) + self.position_embedding = PositionalEmbedding(d_model=d_model) + self.temporal_embedding = nn.Linear(temporal_size, d_model) + self.seqid_embedding = nn.Embedding(seq_num, d_model) + + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + x = self.value_embedding(x) + self.position_embedding(x) + self.temporal_embedding(x_mark[:, :, :-1])\ + + self.seqid_embedding(x_mark[:, :, -1].long()) + + return self.dropout(x) + +"""The SingleStepEmbedding is used by all datasets for single step forecasting.""" +class SingleStepEmbedding(nn.Module): + def __init__(self, cov_size, num_seq, d_model, input_size, device): + super().__init__() + + self.cov_size = cov_size + self.num_class = num_seq + self.cov_emb = nn.Linear(cov_size+1, d_model) + padding = 1 if torch.__version__>='1.5.0' else 2 + self.data_emb = nn.Conv1d(in_channels=1, out_channels=d_model, kernel_size=3, padding=padding, padding_mode='circular') + + self.position = torch.arange(input_size, device=device).unsqueeze(0) + self.position_vec = torch.tensor([math.pow(10000.0, 2.0 * (i // 2) / d_model) for i in range(d_model)], device=device) + + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_(m.weight,mode='fan_in',nonlinearity='leaky_relu') + elif isinstance(m, nn.Linear): + nn.init.xavier_normal_(m.weight) + nn.init.constant_(m.bias, 0) + + def transformer_embedding(self, position, vector): + """ + Input: batch*seq_len. + Output: batch*seq_len*d_model. + """ + result = position.unsqueeze(-1) / vector + result[:, :, 0::2] = torch.sin(result[:, :, 0::2]) + result[:, :, 1::2] = torch.cos(result[:, :, 1::2]) + return result + + def forward(self, x): + covs = x[:, :, 1:(1+self.cov_size)] + seq_ids = ((x[:, :, -1] / self.num_class) - 0.5).unsqueeze(2) + covs = torch.cat([covs, seq_ids], dim=-1) + cov_embedding = self.cov_emb(covs) + data_embedding = self.data_emb(x[:, :, 0].unsqueeze(2).permute(0, 2, 1)).transpose(1,2) + embedding = cov_embedding + data_embedding + + position = self.position.repeat(len(x), 1).to(x.device) + position_emb = self.transformer_embedding(position, self.position_vec.to(x.device)) + + embedding += position_emb + + return embedding diff --git a/Pyraformer/pyraformer/graph_attention.py b/Pyraformer/pyraformer/graph_attention.py new file mode 100644 index 0000000..754cc31 --- /dev/null +++ b/Pyraformer/pyraformer/graph_attention.py @@ -0,0 +1,580 @@ +""" +Test the time and CUDA memory consumption of different attention mechanisms. +""" + +from typing import List +import math +import torch +from torch import nn +import torch.nn.functional as F +import torch.optim as optim +from hierarchical_mm_tvm import graph_mm as graph_mm_tvm +import argparse +import time +import numpy as np +from math import sqrt + +torch.cuda.set_device(0) +print('Using device: {}'.format(torch.cuda.get_device_name())) +import pynvml +pynvml.nvmlInit() + + +def get_q_k(input_size, window_size, stride, device): + """Get the query-key index for PAM-TVM""" + second_length = input_size // stride + second_last = input_size - (second_length - 1) * stride + third_start = input_size + second_length + third_length = second_length // stride + third_last = second_length - (third_length - 1) * stride + max_attn = max(second_last, third_last) + fourth_start = third_start + third_length + fourth_length = third_length // stride + full_length = fourth_start + fourth_length + fourth_last = third_length - (fourth_length - 1) * stride + max_attn = max(third_last, fourth_last) + + max_attn += window_size + 1 + mask = torch.zeros(full_length, max_attn, dtype=torch.int32, device=device) - 1 + + # 按照层内、下层、上层的顺序为序列中每个q找对应的k + # 第一层 + for i in range(input_size): + mask[i, 0:window_size] = i + torch.arange(window_size) - window_size // 2 + # 当window在序列右端时，把它给注释掉 + mask[i, mask[i] > input_size - 1] = -1 + + mask[i, -1] = i // stride + input_size + mask[i][mask[i] > third_start - 1] = third_start - 1 + # 第二层 + for i in range(second_length): + mask[input_size+i, 0:window_size] = input_size + i + torch.arange(window_size) - window_size // 2 + # 当window在序列左端时，置为-1 + mask[input_size+i, mask[input_size+i] < input_size] = -1 + # 当window在序列右端时，置为-1 + mask[input_size+i, mask[input_size+i] > third_start - 1] = -1 + + if i < second_length - 1: + mask[input_size+i, window_size:(window_size+stride)] = torch.arange(stride) + i * stride + else: + mask[input_size+i, window_size:(window_size+second_last)] = torch.arange(second_last) + i * stride + + mask[input_size+i, -1] = i // stride + third_start + mask[input_size+i, mask[input_size+i] > fourth_start - 1] = fourth_start - 1 + # 第三层 + for i in range(third_length): + mask[third_start+i, 0:window_size] = third_start + i + torch.arange(window_size) - window_size // 2 + # 当window在序列左端时，置为-1 + mask[third_start+i, mask[third_start+i] < third_start] = -1 + # 当window在序列右端时，置为-1 + mask[third_start+i, mask[third_start+i] > fourth_start - 1] = -1 + + if i < third_length - 1: + mask[third_start+i, window_size:(window_size+stride)] = input_size + torch.arange(stride) + i * stride + else: + mask[third_start+i, window_size:(window_size+third_last)] = input_size + torch.arange(third_last) + i * stride + + mask[third_start+i, -1] = i // stride + fourth_start + mask[third_start+i, mask[third_start+i] > full_length - 1] = full_length - 1 + # 第四层 + for i in range(fourth_length): + mask[fourth_start+i, 0:window_size] = fourth_start + i + torch.arange(window_size) - window_size // 2 + # 当window在序列左端时，置为-1 + mask[fourth_start+i, mask[fourth_start+i] < fourth_start] = -1 + # 当window在序列右端时，置为-1 + mask[fourth_start+i, mask[fourth_start+i] > full_length - 1] = -1 + + if i < fourth_length - 1: + mask[fourth_start+i, window_size:(window_size+stride)] = third_start + torch.arange(stride) + i * stride + else: + mask[fourth_start+i, window_size:(window_size+fourth_last)] = third_start + torch.arange(fourth_last) + i * stride + + return mask + + +def get_k_q(q_k_mask): + """Get the key-query index from query-key index for PAM-TVM""" + k_q_mask = q_k_mask.clone() + for i in range(len(q_k_mask)): + for j in range(len(q_k_mask[0])): + if q_k_mask[i, j] >= 0: + k_q_mask[i, j] = torch.where(q_k_mask[q_k_mask[i, j]] ==i )[0] + + return k_q_mask + + +def get_mask(input_size, window_size, inner_size, device): + """Get the attention mask of PAM-Naive""" + # Get the size of all layers + all_size = [] + all_size.append(input_size) + second_size = math.floor(input_size / window_size) + all_size.append(second_size) + third_size = math.floor(second_size / window_size) + all_size.append(third_size) + fourth_size = math.floor(third_size / window_size) + all_size.append(fourth_size) + + seq_length = sum(all_size) + mask = torch.zeros(seq_length, seq_length, device=device) + + # Get the intra-scale mask of each scale + inner_window = inner_size // 2 + # The first scale + for i in range(input_size): + left_side = max(i - inner_window, 0) + right_side = min(i + inner_window + 1, input_size) + mask[i, left_side:right_side] = 1 + # The second scale + start = input_size + for i in range(start, start + second_size): + left_side = max(i - inner_window, start) + right_side = min(i + inner_window + 1, start + second_size) + mask[i, left_side:right_side] = 1 + # The third scale + start = input_size + second_size + for i in range(start, start + third_size): + left_side = max(i - inner_window, start) + right_side = min(i + inner_window + 1, start + third_size) + mask[i, left_side:right_side] = 1 + # The fourth scale + start = input_size + second_size + third_size + for i in range(start, start + fourth_size): + left_side = max(i - inner_window, start) + right_side = min(i + inner_window + 1, start + fourth_size) + mask[i, left_side:right_side] = 1 + + # Get the inter-scale mask + start = input_size + for i in range(start, start + second_size): + left_side = (i - input_size) * window_size + if i == (start + second_size - 1): + right_side = start + else: + right_side = (i - input_size + 1) * window_size + mask[i, left_side:right_side] = 1 + mask[left_side:right_side, i] = 1 + # The third scale + start = input_size + second_size + for i in range(start, start + third_size): + left_side = input_size + (i - start) * window_size + if i == (start + third_size - 1): + right_side = start + else: + right_side = input_size + (i - start + 1) * window_size + mask[i, left_side:right_side] = 1 + mask[left_side:right_side, i] = 1 + # The fourth scale + start = input_size + second_size + third_size + for i in range(start, start + fourth_size): + left_side = input_size + second_size + (i - start) * window_size + if i == (start + fourth_size - 1): + right_side = start + else: + right_side = input_size + second_size + (i - start + 1) * window_size + mask[i, left_side:right_side] = 1 + mask[left_side:right_side, i] = 1 + + mask = (1 - mask).bool() + + return mask, all_size + + +"""PAM""" +class GraphSelfAttention(nn.Module): + def __init__(self, opt): + super(GraphSelfAttention, self).__init__() + self.normalize_before = opt.normalize_before + self.n_head = opt.n_head + self.d_k = opt.d_k + + self.w_qs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_ks = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_vs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + nn.init.xavier_uniform_(self.w_qs.weight) + nn.init.xavier_uniform_(self.w_ks.weight) + nn.init.xavier_uniform_(self.w_vs.weight) + + self.fc = nn.Linear(opt.d_k * opt.n_head, opt.d_model) + nn.init.xavier_uniform_(self.fc.weight) + + self.layer_norm = nn.LayerNorm(opt.d_model, eps=1e-6) + self.dropout_attn = nn.Dropout(opt.dropout) + self.dropout_fc = nn.Dropout(opt.dropout) + self.seq_len = opt.seq_len + self.window_size = opt.window_size + self.stride_size = opt.stride_size + self.q_k_mask = get_q_k(self.seq_len, self.window_size, self.stride_size, opt.device) + self.k_q_mask = get_k_q(self.q_k_mask) + + + def forward(self, hidden_states): + residual = hidden_states + + hidden_states = hidden_states + bsz, seq_len, _ = hidden_states.size() + + q = hidden_states + if self.normalize_before: + q = self.layer_norm(q) + + q = self.w_qs(q) + k = self.w_ks(hidden_states) + v = self.w_vs(hidden_states) + q /= math.sqrt(self.d_k) + + q = q.view(bsz, seq_len, self.n_head, self.d_k) + k = k.view(bsz, seq_len, self.n_head, self.d_k) + q = q.float().contiguous() + k = k.float().contiguous() + # attn_weights.size(): (batch_size, L, num_heads, 11) 另外注意这里设置is_t1_diagonaled为False，用于q和k attention + attn_weights = graph_mm_tvm(q, k, self.q_k_mask, self.k_q_mask, False, 0) + attn_weights = self.dropout_attn(F.softmax(attn_weights, dim=-1)) + + v = v.view(bsz, seq_len, self.n_head, self.d_k) + v = v.float().contiguous() + # 这里用于attention scores和v相乘，注意is_t1_diagonaled=True + attn = graph_mm_tvm(attn_weights, v, self.q_k_mask, self.k_q_mask, True, 0) + attn = attn.reshape(bsz, seq_len, self.n_head * self.d_k).contiguous() + context = self.dropout_fc(self.fc(attn)) + context += residual + + if not self.normalize_before: + context = self.layer_norm(context) + + return context + + +"""Multi-head self attention""" +class NormalSelfAttention(nn.Module): + def __init__(self, opt): + super(NormalSelfAttention, self).__init__() + self.normalize_before = opt.normalize_before + self.n_head = opt.n_head + self.d_k = opt.d_k + + self.w_qs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_ks = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_vs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + nn.init.xavier_uniform_(self.w_qs.weight) + nn.init.xavier_uniform_(self.w_ks.weight) + nn.init.xavier_uniform_(self.w_vs.weight) + + self.fc = nn.Linear(opt.d_k * opt.n_head, opt.d_model) + nn.init.xavier_uniform_(self.fc.weight) + + self.layer_norm = nn.LayerNorm(opt.d_model, eps=1e-6) + self.dropout_attn = nn.Dropout(opt.dropout) + self.dropout_fc = nn.Dropout(opt.dropout) + self.seq_len = opt.seq_len + self.window_size = opt.window_size + self.stride_size = opt.stride_size + if opt.mask: + self.mask, _ = get_mask(self.seq_len, self.stride_size, self.window_size, opt.device) + else: + self.mask = None + + + def forward(self, hidden_states): + residual = hidden_states + + hidden_states = hidden_states + bsz, seq_len, _ = hidden_states.size() + + q = hidden_states + if self.normalize_before: + q = self.layer_norm(q) + + q = self.w_qs(q) + k = self.w_ks(hidden_states) + v = self.w_vs(hidden_states) + q /= math.sqrt(self.d_k) + + q = q.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + k = k.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + q = q.float().contiguous() + k = k.float().contiguous() + v = v.float().contiguous() + + attn = torch.matmul(q, k.transpose(2, 3)) + + if self.mask is not None: + attn = attn.masked_fill(self.mask.unsqueeze(0).unsqueeze(1), -1e9) + + attn = self.dropout_attn(F.softmax(attn, dim=-1)) + attn = torch.matmul(attn, v).transpose(1, 2).contiguous() + attn = attn.view(bsz, seq_len, self.n_head * self.d_k) + + context = self.dropout_fc(self.fc(attn)) + context += residual + + if not self.normalize_before: + context = self.layer_norm(context) + + return context + + +"""Prob-sparse attention""" +class ProbSparseAttention(nn.Module): + def __init__(self, opt): + super(ProbSparseAttention, self).__init__() + self.normalize_before = opt.normalize_before + self.n_head = opt.n_head + self.d_k = opt.d_k + + self.w_qs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_ks = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + self.w_vs = nn.Linear(opt.d_model, opt.n_head * opt.d_k, bias=False) + nn.init.xavier_uniform_(self.w_qs.weight) + nn.init.xavier_uniform_(self.w_ks.weight) + nn.init.xavier_uniform_(self.w_vs.weight) + + self.fc = nn.Linear(opt.d_k * opt.n_head, opt.d_model) + nn.init.xavier_uniform_(self.fc.weight) + + self.layer_norm = nn.LayerNorm(opt.d_model, eps=1e-6) + self.dropout_attn = nn.Dropout(opt.dropout) + self.dropout_fc = nn.Dropout(opt.dropout) + self.seq_len = opt.seq_len + self.factor = opt.factor + + def _prob_QK(self, Q, K, sample_k, n_top): # n_top: c*ln(L_q) + # Q [B, H, L, D] + B, H, L_K, E = K.shape + _, _, L_Q, _ = Q.shape + + # calculate the sampled Q_K + K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E) + index_sample = torch.randint(L_K, (L_Q, sample_k)) # real U = U_part(factor*ln(L_k))*L_q + K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :] + Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze() + + # find the Top_k query with sparisty measurement + M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K) + M_top = M.topk(n_top, sorted=False)[1] + + # use the reduced Q to calculate Q_K + Q_reduce = Q[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + M_top, :] # factor*ln(L_q) + Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) # factor*ln(L_q)*L_k + + return Q_K, M_top + + def _get_initial_context(self, V, L_Q): + B, H, L_V, D = V.shape + V_sum = V.mean(dim=-2) + contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone() + + return contex + + def _update_context(self, context_in, V, scores, index, L_Q): + B, H, L_V, D = V.shape + + attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores) + + context_in[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + index, :] = torch.matmul(attn, V).type_as(context_in) + return context_in + + def forward(self, hidden_states): + residual = hidden_states + + hidden_states = hidden_states + bsz, seq_len, _ = hidden_states.size() + + q = hidden_states + if self.normalize_before: + q = self.layer_norm(q) + + q = self.w_qs(q) + k = self.w_ks(hidden_states) + v = self.w_vs(hidden_states) + q /= math.sqrt(self.d_k) + + q = q.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + k = k.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + v = v.view(bsz, seq_len, self.n_head, self.d_k).transpose(1, 2) + q = q.float().contiguous() + k = k.float().contiguous() + v = v.float().contiguous() + + u = U_part = self.factor * np.ceil(np.log(seq_len)).astype('int').item() # c*ln(L_k) + + U_part = U_part if U_part