Adding pyraformer(not working)

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",
+]

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)
+
+        

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)
+
+
+
+

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,
+        )

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
+

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
+

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
+

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
+

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

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<seq_len else seq_len
+        u = u if u < seq_len else seq_len
+
+        scores_top, index = self._prob_QK(q, k, sample_k=U_part, n_top=u) 
+
+        # get the context
+        context = self._get_initial_context(v, seq_len)
+        # update the context with selected top_k queries
+        context = self._update_context(context, v, scores_top, index, seq_len).transpose(1, 2).contiguous()
+
+        context = context.view(bsz, seq_len, self.n_head * self.d_k)
+
+        context = self.dropout_fc(self.fc(context))
+        context += residual
+
+        if not self.normalize_before:
+            context = self.layer_norm(context)
+
+        return context
+
+
+def parsing():
+    parser = argparse.ArgumentParser(description='Needed for graph self attention.')
+    parser.add_argument('-d_model', type=int, default=256)
+    parser.add_argument('-d_k', type=int, default=64)
+    parser.add_argument('-normalize_before', type=bool, default=False)
+    parser.add_argument('-n_head', type=int, default=4)
+    parser.add_argument('-dropout', type=float, default=0.1)
+
+    # arguments for Multiformer
+    parser.add_argument('-window_size', type=int, default=3)
+    parser.add_argument('-stride_size', type=int, default=25)
+
+    # arguments for ProbSparse
+    parser.add_argument('-factor', type=int, default=5)
+
+    # arguments for full-attention
+    parser.add_argument('-mask', type=int, default=0)
+
+    parser.add_argument('-seq_len', type=int, default=1000)
+    args = parser.parse_args()
+
+    return args
+
+
+def test_NSA(args, input_len):
+    """Test the time and CUDA memory consumption of normal self attention."""
+    handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+    init_mem = meminfo.used / 1024**3
+
+    NSA_Layer = NormalSelfAttention(args).to(args.device)
+    optimizer = optim.Adam(NSA_Layer.parameters(), 1e-4)
+    optimizer.zero_grad()
+    hidden_state = torch.ones(4, input_len, args.d_model, dtype=torch.float32).to(args.device)
+    fake_gt = torch.zeros(4, input_len, args.d_model).to(args.device)
+
+    # Preload the layer
+    result = NSA_Layer(hidden_state)
+    loss = ((fake_gt  - result) ** 2).mean()
+    loss.backward()
+    optimizer.step()
+
+    used_memory = 0
+    start_time = time.time()
+    for i in range(1000):
+        result = NSA_Layer(hidden_state)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        used_memory += meminfo.used / 1024**3
+        loss = ((fake_gt  - result) ** 2).mean()
+        loss.backward()
+        optimizer.step()
+
+    print('NSA used average time: {} s'.format(round((time.time() - start_time) / 1000, 4)))
+    used_memory = used_memory / 1000
+    print('NSA used average memory: {} GB'.format(round(used_memory-init_mem, 4)))
+
+
+def test_GSA(args, input_len):
+    """Test the time and CUDA memory consumption of PAM."""
+    handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+    init_mem = meminfo.used / 1024**3
+
+    GSA_Layer = GraphSelfAttention(args).to(args.device)
+    optimizer = optim.Adam(GSA_Layer.parameters(), 1e-4)
+    optimizer.zero_grad()
+    hidden_state = torch.ones(4, input_len, args.d_model, dtype=torch.float32, device=args.device)
+    fake_gt = torch.zeros(4, input_len, args.d_model, device=args.device)
+
+    # Preload the layer
+    result = GSA_Layer(hidden_state)
+    loss = ((fake_gt  - result) ** 2).mean()
+    loss.backward()
+    optimizer.step()
+
+    used_memory = 0
+    repeat_times = 1000
+    start_time = time.time()
+    for i in range(repeat_times):
+        result = GSA_Layer(hidden_state)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        used_memory += meminfo.used / 1024**3
+        loss = ((fake_gt  - result) ** 2).mean()
+        loss.backward()
+        optimizer.step()
+
+    print('GSA used time:{} s'.format(round((time.time() - start_time) / repeat_times, 4)))
+    used_memory = used_memory / repeat_times
+    print('GSA used average memory: {} GB'.format(round(used_memory-init_mem, 4)))
+
+
+def test_PSA(args, input_len):
+    """Test the time and CUDA memory consumption of Prob-sparse self attention."""
+    handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+    init_mem = meminfo.used / 1024**3
+
+    LSA_Layer = ProbSparseAttention(args).to(args.device)
+    optimizer = optim.Adam(LSA_Layer.parameters(), 1e-4)
+    optimizer.zero_grad()
+    hidden_state = torch.ones(4, input_len, args.d_model, dtype=torch.float32, device=args.device)
+    fake_gt = torch.zeros(4, input_len, args.d_model, device=args.device)
+
+    # Preload the layer
+    result = LSA_Layer(hidden_state)
+    loss = ((fake_gt  - result) ** 2).mean()
+    loss.backward()
+    optimizer.step()
+
+    used_memory = 0
+    repeat_times = 1000
+    start_time = time.time()
+    for i in range(repeat_times):
+        result = LSA_Layer(hidden_state)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(1)
+        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        used_memory += meminfo.used / 1024**3
+        loss = ((fake_gt  - result) ** 2).mean()
+        loss.backward()
+        optimizer.step()
+
+    print('LSA used time:{} s'.format(round((time.time() - start_time) / repeat_times, 4)))
+    used_memory = used_memory / repeat_times
+    print('LSA used average memory: {} GB'.format(round(used_memory-init_mem, 4)))
+
+
+if __name__ == '__main__':
+    args = parsing()
+    if torch.cuda.is_available():
+        args.device = torch.device('cuda')
+    else:
+        args.device = torch.device('cpu')
+
+    input_size = args.seq_len
+    stride = args.stride_size
+    second_length = input_size // stride
+    third_length = second_length // stride
+    fourth_length = third_length // stride
+    input_len = input_size + second_length + third_length + fourth_length
+
+    if args.mask:
+        print('sequence length: {}'.format(input_len))
+        test_NSA(args, input_len)
+    else:
+        print('sequence length: {}'.format(input_size))
+        test_NSA(args, input_size)
+
+    print('sequence length: {}'.format(input_len))
+    test_GSA(args, input_len)
+    print('sequence length: {}'.format(input_size))
+    test_PSA(args, input_size)
+

Pyraformer/pyraformer/hierarchical_mm_tvm.py

@@ -0,0 +1,282 @@
+"""
+Modified based on Longformer.
+@article{Beltagy2020Longformer,
+  title={Longformer: The Long-Document Transformer},
+  author={Iz Beltagy and Matthew E. Peters and Arman Cohan},
+  journal={arXiv:2004.05150},
+  year={2020},
+}
+"""
+
+from typing import Union
+from functools import lru_cache
+
+import torch
+import os.path
+import sys
+sys.path.append('pyraformer/tvm/python')
+
+class GraphMM(torch.autograd.Function):
+    '''Class to encapsulate tvm code for compiling a diagonal_mm function, in addition to calling
+    this function from PyTorch
+    '''
+
+    function_dict = {}  # save a list of functions, each has a different set of parameters
+
+    @staticmethod
+    def _compile_function(dtype: str, device: str, b0: int = 4, b1: int = 8, b2: int = 8):
+        '''Compiles a tvm function that computes diagonal_mm
+        args:
+        dtype: str in ['float64', 'float32', 'float16']
+        device: str in ['cpu' or 'cuda']
+        b0, b1, b2: size of tensor tiles. Very important for good performance
+        '''
+        import tvm  # import the full tvm library here for compilation. Don't import at the top of the file in case we don't need to compile
+        from tvm.contrib import nvcc
+        @tvm.register_func
+        def tvm_callback_cuda_compile(code):
+            """Use nvcc compiler for better perf."""
+            ptx = nvcc.compile_cuda(code, target="ptx", arch='sm_52')  # use old arch for this to work on old GPUs
+            return ptx
+
+        assert dtype in ['float16', 'float32', 'float64']
+        assert device in ['cpu', 'cuda']
+        device = None if device == 'cpu' else device
+        tgt_host="llvm"
+
+        b = tvm.te.var('b')  # batch size
+        n = tvm.te.var('n')  # sequence length
+        h = tvm.te.var('h')  # number of heads
+        m = tvm.te.var('m')  # hidden dimension
+        w = tvm.te.var('w')  # window size
+        padding = tvm.te.var('padding')  # padding
+        transpose_t1 = tvm.te.var('transpose_t1')  # t1 should be transposed
+        t1d3 = tvm.te.var('t1d3')  # last dimension of t1
+        t3d3 = tvm.te.var('t3d3')  # last dimension of t3 (the result tensor)
+        max_attn = tvm.te.var('max_attn')
+        X = tvm.te.placeholder((b, n, h, t1d3), name='X', dtype=dtype)  # first tensor
+        Y = tvm.te.placeholder((b, n, h, m), name='Y', dtype=dtype)  # second tensor
+        k = tvm.te.reduce_axis((0, t1d3), name='k')  # dimension to sum over
+        q_k_mask = tvm.te.placeholder((n, max_attn), name='q_k', dtype='int')  # dilation per head
+        k_q_mask = tvm.te.placeholder((n, max_attn), name='k_q', dtype='int') # 
+        output_shape = (b, n, h, t3d3)  # shape of the result tensor
+
+        algorithm = lambda l, i, q, j: tvm.te.sum(
+            tvm.te.if_then_else(
+                t3d3 == m,  # if output dimension == m, then t1 is diagonaled (FIXME: This breaks if t3d3 == m == t1d3)
+                tvm.te.if_then_else(
+                    transpose_t1 == 0,
+                    tvm.te.if_then_else(
+                        q_k_mask[i, k]>=0,
+                        X[l, i, q, k] * Y[l, q_k_mask[i, k], q, j],  # t1 is diagonaled
+                        padding
+                    ),
+                    tvm.te.if_then_else(
+                        q_k_mask[i, k]>=0,
+                        X[l, q_k_mask[i, k], q, k_q_mask[i, k]] * Y[l, q_k_mask[i, k], q, j],  # # t1 is diagonaled and should be transposed
+                        padding
+                    ),
+                ),
+                tvm.te.if_then_else(
+                    q_k_mask[i, j]>=0,
+                    X[l, i, q, k] * Y[l, q_k_mask[i, j], q, k],  # t1 is not diagonaled, but the output tensor is going to be
+                    padding
+                )
+            ), axis=k)
+
+        Z = tvm.te.compute(output_shape, algorithm, name='Z')  # automatically generate cuda code
+        s = tvm.te.create_schedule(Z.op)
+
+        print('Lowering: \n ===================== \n{}'.format(tvm.lower(s, [X, Y, q_k_mask, k_q_mask], simple_mode=True)))
+
+        # split long axis into smaller chunks and assing each one to a separate GPU thread/block
+        ko, ki = s[Z].split(Z.op.reduce_axis[0], factor=b0)
+        ZF = s.rfactor(Z, ki)
+
+        j_outer, j_inner = s[Z].split(s[Z].op.axis[-1], factor=b1)
+        i_outer, i_inner = s[Z].split(s[Z].op.axis[1], factor=b2)
+
+        s[Z].bind(j_outer, tvm.te.thread_axis("blockIdx.x"))
+        s[Z].bind(j_inner, tvm.te.thread_axis("threadIdx.y"))
+
+        s[Z].bind(i_outer, tvm.te.thread_axis("blockIdx.y"))
+        s[Z].bind(i_inner, tvm.te.thread_axis("threadIdx.z"))
+
+        tx = tvm.te.thread_axis("threadIdx.x")
+        s[Z].bind(s[Z].op.reduce_axis[0], tx)
+        s[ZF].compute_at(s[Z], s[Z].op.reduce_axis[0])
+        s[Z].set_store_predicate(tx.var.equal(0))
+
+        print('Lowering with GPU splits: \n ===================== \n{}'.format(tvm.lower(s, [X, Y, q_k_mask, k_q_mask], simple_mode=True)))
+
+        # compiling the automatically generated cuda code
+        graph_mm = tvm.build(s, [X, Y, Z, q_k_mask, k_q_mask, max_attn, padding, transpose_t1, t3d3], target=device, target_host=tgt_host, name='graph_mm')
+        return graph_mm
+
+    @staticmethod
+    def _get_lib_filename(dtype: str, device: str):
+        base_filename = 'lib/lib_hierarchical_mm'
+        return '{}_{}_{}.so'.format(base_filename, dtype, device)
+
+    @staticmethod
+    def _save_compiled_function(f, dtype: str, device: str):
+        if not os.path.exists('lib/'):
+            os.makedirs('lib/')
+        f.export_library(GraphMM._get_lib_filename(dtype, device))
+
+    @staticmethod
+    def _load_compiled_function(dtype: str, device: str):
+        # from tvm.module import load  # this can be the small runtime python library, and doesn't need to be the whole thing
+        from tvm.runtime.module import load_module as load
+
+        filename = GraphMM._get_lib_filename(dtype, device)
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        potential_dirs = ['../../', '../', './', f'{current_dir}/', f'{current_dir}/../']
+        for potential_dir in  potential_dirs:
+            filepath = '{}{}'.format(potential_dir, filename)
+            if os.path.isfile(filepath):
+                print('Loading tvm binary from: {}'.format(filepath))
+                return load(filepath)
+        return None
+
+    @staticmethod
+    def _get_function(dtype: str, device: str):
+        '''Loads the function from the disk or compile it'''
+        # A list of arguments that define the function
+        args = (dtype, device)
+        if args not in GraphMM.function_dict:
+            graph_mm = GraphMM._load_compiled_function(dtype, device)  # try to load from disk
+            if not graph_mm:
+                print('Tvm binary not found. Compiling ...')
+                graph_mm = GraphMM._compile_function(dtype, device)  # compile
+                GraphMM._save_compiled_function(graph_mm, dtype, device)  # save to disk
+            # convert the tvm function into a pytorch function
+            from tvm.contrib import dlpack
+            graph_mm_pytorch = dlpack.to_pytorch_func(graph_mm)  # wrap it as a pytorch function
+            # save the function into a dictionary to be reused
+            GraphMM.function_dict[args] = graph_mm_pytorch  # save it in a dictionary for next time
+        return GraphMM.function_dict[args]
+
+    @staticmethod
+    def _graph_mm(t1: torch.Tensor, t2: torch.Tensor, q_k_mask: torch.Tensor, k_q_mask: torch.Tensor,
+                       is_t1_diagonaled: bool = False, transpose_t1: bool = False, padding: int = 0,
+                       autoregressive: bool = False):
+        '''Calls the compiled function after checking the input format. This function is called in three different modes.
+        t1 x t2 = r ==> t1 and t2 are not diagonaled, but r is. Useful for query x key = attention_scores
+        t1 x t2 = r ==> t1 is diagonaled, but t2 and r are not. Useful to compuate attantion_scores x value = context
+        t1 x t2 = r ==> t1 is diagonaled and it should be transposed, but t2 and r are not diagonaled. Useful in some of
+                            the calculations in the backward pass.
+        '''
+        dtype = str(t1.dtype).split('.')[1]
+        device = t1.device.type
+        assert len(t1.shape) == 4
+        assert len(t1.shape) == len(t2.shape)
+        assert t1.shape[:3] == t2.shape[:3]
+
+        b = t1.shape[0]  # batch size
+        n = t1.shape[1]  # sequence length
+        h = t1.shape[2]  # number of heads
+        m = t2.shape[3]  # hidden dimension
+        max_attn = q_k_mask.size(1)
+        if is_t1_diagonaled:
+            assert t1.shape[3] == max_attn
+            r = t1.new_empty(b, n, h, m)  # allocate spase for the result tensor
+        else:
+            assert not transpose_t1
+            assert t1.shape[3] == m
+            r = t1.new_empty(b, n, h, max_attn)  # allocate spase for the result tensor
+
+        # gets function from memory, from disk or compiles it from scratch
+        _graph_mm_function = GraphMM._get_function(dtype=dtype, device=device)
+
+        # The last argument to this function is a little hacky. It is the size of the last dimension of the result tensor
+        # We use it as a proxy to tell if t1_is_diagonaled or not (if t1 is diagonaled, result is not, and vice versa).
+        # The second reason is that the lambda expression in `_compile_function` is easier to express when the shape
+        # of the output is known
+        # This functions computes diagonal_mm then saves the result in `r`
+        if m == max_attn:
+            # FIXME
+            print('Error: the hidden dimension {m} shouldn\'t match number of diagonals {c}')
+            assert False
+        _graph_mm_function(t1, t2, r, q_k_mask, k_q_mask, max_attn, padding, transpose_t1, m if is_t1_diagonaled else max_attn)
+        return r
+
+    @staticmethod
+    def _prepare_tensors(t):
+        '''Fix `stride()` information of input tensor. This addresses some inconsistency in stride information in PyTorch.
+        For a tensor t, if t.size(0) == 1, then the value of t.stride()[0] doesn't matter.
+        TVM expects this value to be the `product(t.size()[1:])` but PyTorch some times sets it to `t.stride()[1]`.
+        Here's an example to reporduce this issue:
+            import torch
+            print(torch.randn(1, 10).stride())
+            > (10, 1)
+            print(torch.randn(10, 1).t().contiguous().stride())
+            > (1, 1)  # expected it to be (10, 1) as above
+            print(torch.randn(10, 2).t().contiguous().stride())
+            > (10, 1) # but gets the expected stride if the first dimension is > 1
+        '''
+        assert t.is_contiguous()
+        t_stride = list(t.stride())
+        t_size = list(t.size())
+        # Fix wrong stride information for the first dimension. This occures when batch_size=1
+        if t_size[0] == 1 and t_stride[0] == t_stride[1]:
+            # In this case, the stride of the first dimension should be the product
+            # of the sizes  of all other dimensions
+            t_stride[0] = t_size[1] * t_size[2] * t_size[3]
+            t = t.as_strided(size=t_size, stride=t_stride)
+        return t
+
+    min_seq_len = 16  # unexpected output if seq_len < 16
+
+    @staticmethod
+    def forward(ctx, t1: torch.Tensor, t2: torch.Tensor, q_k_mask, k_q_mask, is_t1_diagonaled: bool = False, padding: int = 0) -> torch.Tensor:
+        '''Compuates diagonal_mm of t1 and t2.
+        args: 
+        t1: torch.Tensor = (batch_size, seq_len, num_attention_heads, hidden_size|number_of_diagonals).
+            t1 can be a regular tensor (e.g. `query_layer`) or a diagonaled one (e.g. `attention_scores`)
+        t2: torch.Tensor = (batch_size, seq_len, num_attention_heads, hidden_size). This is always a non-diagonaled
+            tensor, e.g. `key_layer` or `value_layer`
+        w: int = window size; number of attentions on each side of the word
+        d: torch.Tensor or int = dilation of attentions per attention head. If int, the same dilation value will be used for all
+            heads. If torch.Tensor, it should be 1D of lenth=number of attention heads
+        is_t1_diagonaled: is t1 a diagonaled or a regular tensor
+        padding: the padding value to use when accessing invalid locations. This is mainly useful when the padding
+            needs to be a very large negative value (to compute softmax of attentions). For other usecases,
+            please use zero padding.
+        autoregressive: if true, return only the lower triangle
+        returns: torch.Tensor = (batch_size, seq_len, num_attention_heads, hidden_size|number_of_diagonals)
+            if t1 is diagonaed, result is non-diagonaled, and vice versa
+        '''
+        seq_len = t1.size(1)
+        assert seq_len >= GraphMM.min_seq_len, 'avoid splitting errors by using seq_len >= {}'.format(GraphMM.min_seq_len)  # FIXME
+
+        t1 = GraphMM._prepare_tensors(t1)
+        t2 = GraphMM._prepare_tensors(t2)
+        q_k_mask = GraphMM._prepare_tensors(q_k_mask)
+        k_q_mask = GraphMM._prepare_tensors(k_q_mask)
+        ctx.save_for_backward(t1, t2, q_k_mask, k_q_mask)
+        ctx.is_t1_diagonaled = is_t1_diagonaled
+        # output = t1.mm(t2)  # what would have been called if this was a regular matmul
+        output = GraphMM._graph_mm(t1, t2, q_k_mask, k_q_mask, is_t1_diagonaled=is_t1_diagonaled, padding=padding)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        t1, t2, q_k_mask, k_q_mask = ctx.saved_tensors
+        is_t1_diagonaled = ctx.is_t1_diagonaled
+        if not grad_output.is_contiguous():
+            grad_output = grad_output.contiguous()  # tvm requires all input tensors to be contiguous
+        grad_output = GraphMM._prepare_tensors(grad_output)
+        # http://cs231n.github.io/optimization-2/
+        # https://pytorch.org/docs/master/notes/extending.html
+        # grad_t1 = grad_output.mm(t2)  # what would have been called if this was a regular matmul
+        grad_t1 = GraphMM._graph_mm(grad_output, t2, q_k_mask, k_q_mask, is_t1_diagonaled=not is_t1_diagonaled)
+        # grad_t2 = grad_output.t().mm(t1)  # or `grad_t2 = t1.t().mm(grad_output).t()` because `(AB)^T = B^TA^T`
+        if is_t1_diagonaled:
+            grad_t2 = GraphMM._graph_mm(t1, grad_output, q_k_mask, k_q_mask, is_t1_diagonaled=True, transpose_t1=True)
+        else:
+            grad_t2 = GraphMM._graph_mm(grad_output, t1, q_k_mask, k_q_mask, is_t1_diagonaled=True, transpose_t1=True)
+        return grad_t1, grad_t2, None, None, None, None, None
+
+
+graph_mm = GraphMM.apply

Pyraformer/tools.py

@@ -0,0 +1,93 @@
+from torch.nn.modules import loss
+import torch
+import numpy as np
+
+
+def MAE(pred, true):
+    return np.mean(np.abs(pred-true))
+
+def MSE(pred, true):
+    return np.mean((pred-true)**2)
+
+def RMSE(pred, true):
+    return np.sqrt(MSE(pred, true))
+
+def MAPE(pred, true):
+    return np.mean(np.abs((pred - true) / true))
+
+def MSPE(pred, true):
+    return np.mean(np.square((pred - true) / true))
+
+def metric(pred, true):
+    mae = MAE(pred, true)
+    mse = MSE(pred, true)
+    rmse = RMSE(pred, true)
+    mape = MAPE(pred, true)
+    mspe = MSPE(pred, true)
+
+    return mae,mse,rmse,mape,mspe
+
+class StandardScaler():
+    def __init__(self):
+        self.mean = 0.
+        self.std = 1.
+
+    def fit(self, data):
+        self.mean = data.mean(0)
+        self.std = data.std(0)
+
+    def transform(self, data):
+        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
+        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
+        return (data - mean) / std
+
+    def inverse_transform(self, data):
+        mean = torch.from_numpy(self.mean).type_as(data).to(data.device) if torch.is_tensor(data) else self.mean
+        std = torch.from_numpy(self.std).type_as(data).to(data.device) if torch.is_tensor(data) else self.std
+        return (data * std) + mean
+
+class TopkMSELoss(torch.nn.Module):
+    def __init__(self, topk) -> None:
+        super().__init__()
+        self.topk = topk
+        self.criterion = torch.nn.MSELoss(reduction='none')
+
+    def forward(self, output, label):
+        losses = self.criterion(output, label).mean(2).mean(1)
+        losses = torch.topk(losses, self.topk)[0]
+
+        return losses
+
+class SingleStepLoss(torch.nn.Module):
+    """ Compute top-k log-likelihood and mse. """
+
+    def __init__(self, ignore_zero):
+        super().__init__()
+        self.ignore_zero = ignore_zero
+
+    def forward(self, mu, sigma, labels, topk=0):
+        if self.ignore_zero:
+            indexes = (labels != 0)
+        else:
+            indexes = (labels >= 0)
+
+        distribution = torch.distributions.normal.Normal(mu[indexes], sigma[indexes])
+        likelihood = -distribution.log_prob(labels[indexes])
+
+        diff = labels[indexes] - mu[indexes]
+        se = diff * diff
+
+        if 0 < topk < len(likelihood):
+            likelihood = torch.topk(likelihood, topk)[0]
+            se = torch.topk(se, topk)[0]
+
+        return likelihood, se
+
+def AE_loss(mu, labels, ignore_zero):
+    if ignore_zero:
+        indexes = (labels != 0)
+    else:
+        indexes = (labels >= 0)
+
+    ae = torch.abs(labels[indexes] - mu[indexes])
+    return ae