+
+
> Multivariate Time Series (MTS) forecasting plays a vital role in a wide range of applications. Recently, Spatial-Temporal Graph Neural Networks (STGNNs) have become increasingly popular MTS forecasting methods. STGNNs jointly model the spatial and temporal patterns of MTS through graph neural networks and sequential models, significantly improving the prediction accuracy. But limited by model complexity, most STGNNs only consider short-term historical MTS data, such as data over the past one hour. However, the patterns of time series and the dependencies between them (i.e., the temporal and spatial patterns) need to be analyzed based on long-term historical MTS data. To address this issue, we propose a novel framework, in which STGNN is Enhanced by a scalable time series Pre-training model (STEP). Specifically, we design a pre-training model to efficiently learn temporal patterns from very long-term history time series (e.g., the past two weeks) and generate segment-level representations. These representations provide contextual information for short-term time series input to STGNNs and facilitate modeling dependencies between time series. Experiments on three public real-world datasets demonstrate that our framework is capable of significantly enhancing downstream STGNNs, and our pre-training model aptly captures temporal patterns.
-## 1. Table of Contents
+## ๐ Table of Contents
```text
-config --> Training configs and model configs for each dataset
-dataloader --> MTS dataset
-easytorch --> EasyTorch
-model --> Model architecture
-checkpoints --> Saving the checkpoints according to md5 of the configuration file
-datasets --> Raw datasets and preprocessed data
-train_logs --> Our train logs.
-TSFormer_CKPT --> Our checkpoints.
-```
+basicts --> The BasicTS, which provides standard pipelines for training MTS forecasting models. Don't worry if you don't know it, because it doesn't prevent you from understanding STEP's code.
-## 2. Requirements
+datasets --> Raw datasets and preprocessed data
-```bash
-pip install -r requirements.txt
-```
+figures --> Some figures used in README.
-## 3. Data Preparation
+scripts --> Data preprocessing scripts.
-### 3.1 Download Data
+step --> The implementation of STEP, including the architecture, dataloader, loss, and runner for STEP.
-Download data from link [Google Drive](https://drive.google.com/drive/folders/1F7fEdXpnEQ75sxQval52jN4r3ZKoufGV?usp=sharing) or [BaiduYun](https://pan.baidu.com/s/1IWVhsvKpxHEb2CFLCR7P3A?pwd=w1wJ) to the code root directory.
+tsformer_ckpt --> Pre-trained TSFormer for METR-LA, PEMS-BAY, and PEMS04 dataset.
+```
-Then, unzip data by:
+## ๐ฟ Requirements
```bash
-unzip TSFormer_CKPT.zip
-mkdir datasets
-unzip raw_data.zip -d datasets
-unzip sensor_graph.zip -d datasets
-rm *.zip
+pip install -r requirements.txt
```
-`TSFormer_CKPT/` contains the pre-trained model for each dataset.
-You can also find all the training logs of the pre-training stage and forecasting stage in `training_logs/`.
+## ๐ฆ Data Preparation
-### 3.2 Preprocess Data
+### **Download Raw Data**
-```bash
-python datasets/raw_data/$DATASET_NAME/generate_data.py
-```
-
-Replace `$DATASET_NAME` with one of `METR-LA`, `PEMS-BAY`, `PEMS04`.
+You can download all the raw datasets at [Google Drive](https://drive.google.com/file/d/1PY7IZ3SchpyXfNIXs71A2GEV29W5QCv2/view?usp=sharing) or [Baidu Yun](https://pan.baidu.com/s/1CXLxeHxHIMWLy3IKGFUq8g?pwd=blf8), and unzip them to `datasets/raw_data/`.
-The processed data is placed in `datasets/$DATASET_NAME`.
+### **Pre-process Data**
-## 4. Train STEP based on a Pre-trained TSFormer
+You can pre-process all data via:
```bash
-python main.py --cfg='config/$DATASET/forecasting.py' --gpu='0, 1'
-# python main.py --cfg='config/METR-LA/forecasting.py' --gpu='0, 1'
-# python main.py --cfg='config/PEMS-BAY/forecasting.py' --gpu='0, 1'
-# python main.py --cfg='config/PEMS04/forecasting.py' --gpu='0, 1'
+cd /path/to/your/project
+bash scripts/data_preparation/all.sh
```
-Replace `$DATASET_NAME` with one of `METR-LA`, `PEMS-BAY`, `PEMS04` as shown in the code above.
+Then the `dataset` directory will look like this:
-Configuration file
-`config/$DATASET_NAME/forecasting.py` describes the forecasting configurations.
+```text
+datasets
+ โโMETR-LA
+ โโMETR-BAY
+ โโPEMS04
+ โโraw_data
+ | โโPEMS04
+ | โโPEMS-BAY
+ | โโMETR-LA
+ โโREADME.md
+```
-We use 2 GPU for forecasting stage as default, edit `GPU_NUM` property in the configuration file and `--gpu` in the command line to run on your own hardware.
+## ๐ฏ Train STEP based on a Pre-trained TSFormer
-Note that different GPU numbers lead to different real batch sizes, affecting the learning rate setting and the forecasting accuracy.
+```bash
+python step/run.py --cfg='step/step_$DATASET.py' --gpus='0, 1'
+# python step/run.py --cfg='step/step_METR-LA.py' --gpus='0'
+# python step/run.py --cfg='step/step_PEMS-BAY.py' --gpus='0, 1'
+# python step/run.py --cfg='step/step_PEMS04.py' --gpus='0, 1'
+```
+
+Replace `$DATASET_NAME` with one of `METR-LA`, `PEMS-BAY`, `PEMS04` as shown in the code above.
+Configuration file `step/STEP_$DATASET.py` describes the forecasting configurations.
+Edit `BATCH_SIZE` and `GPU_NUM` in the configuration file and `--gpu` in the command line to run on your own hardware.
+Note that different GPU number leads to different real batch sizes, affecting the learning rate setting and the forecasting accuracy.
Our training logs are shown in `train_logs/Backend_metr.log`, `train_logs/Backend_pems04.log`, and `train_logs/Backend_pemsbay.log`.
-## 5. Train STEP from Scratch
+## โ Train STEP from Scratch
-### 5.1 Pre-training Stage
+### **Pre-training Stage**
```bash
-python main.py --cfg='config/$DATASET/pretraining.py' --gpu='0'
-# python main.py --cfg='config/METR-LA/pretraining.py' --gpu='0'
-# python main.py --cfg='config/PEMS-BAY/pretraining.py' --gpu='0, 1, 2, 3, 4, 5, 6, 7'
-# python main.py --cfg='config/PEMS04/pretraining.py' --gpu='0, 1'
+python step/run.py --cfg='step/TSFormer_$DATASET.py' --gpus '0'
+# python step/run.py --cfg='step/TSFormer_METR-LA.py' --gpus='0'
+# python step/run.py --cfg='step/TSFormer_PEMS-BAY.py' --gpus='0, 1'
+# python step/run.py --cfg='step/TSFormer_PEMS04.py' --gpus='0'
```
Replace `$DATASET_NAME` with one of `METR-LA`, `PEMS-BAY`, `PEMS04` as shown in the code above.
-
-Configuration file `config/$DATASET_NAME/pretraining.py` describes the pre-training configurations.
-
+Configuration file `step/TSFormer_$DATASET.py` describes the pre-training configurations.
Edit the `BATCH_SIZE` and `GPU_NUM` in the configuration file and `--gpu` in the command line to run on your own hardware.
+All the training logs, including the config file, training log, and checkpoints, will be saved in `checkpoints/MODEL_EPOCH/MD5_of_config_file`.
+For example, `checkpoints/TSFormer_100/5afe80b3e7a3dc055158bcfe99afbd7f`.
-### 5.2 Forecasting Stage
+### **Forecasting Stage**
-Move your pre-trained model checkpoints to `TSFormer_CKPT/`.
+After pre-training TSFormer, move your pre-trained best checkpoint to `tsformer_ckpt/`.
For example:
```bash
-cp checkpoints/TSFormer_200/9b4b52e25a30aabd21dc1c9429063196/TSFormer_180.pt TSFormer_CKPT/TSFormer_PEMS-BAY.pt
-```
-
-```bash
-cp checkpoints/TSFormer_200/fac3814778135a6d46063e3cab20257c/TSFormer_147.pt TSFormer_CKPT/TSFormer_PEMS04.pt
+cp checkpoints/TSFormer_100/5afe80b3e7a3dc055158bcfe99afbd7f/TSFormer_best_val_MAE.pt tsformer_ckpt/TSFormer_$DATASET_NAME.pt
```
-```bash
-cp checkpoints/TSFormer_200/3de38a467aef981dd6f24127b6fb5f50/TSFormer_030.pt TSFormer_CKPT/TSFormer_METR-LA.pt
-```
+Replace `$DATASET_NAME` with one of `METR-LA`, `PEMS-BAY`, `PEMS04`.
Then train the downstream STGNN (Graph WaveNet) like in section 4.
-## 6. Performance and Visualization
+## ๐ Performance and Visualization
-## 7. More Related Works
+## ๐ More Related Works
- [D2STGNN: Decoupled Dynamic Spatial-Temporal Graph Neural Network for Traffic Forecasting. VLDB'22.](https://github.com/zezhishao/D2STGNN)
- [BasicTS: An Open Source Standard Time Series Forecasting Benchmark.](https://github.com/zezhishao/BasicTS)
-## 8. Citing
+## Citing
If you find this repository useful for your work, please consider citing it as follows:
diff --git a/Runner_FullModel.py b/Runner_FullModel.py
deleted file mode 100644
index d0f464e..0000000
--- a/Runner_FullModel.py
+++ /dev/null
@@ -1,367 +0,0 @@
-import time
-import math
-
-import torch
-from torch import nn
-from sklearn.metrics import mean_absolute_error
-
-from easytorch import Runner
-from easytorch.utils.dist import master_only
-
-from models.step import STEP
-from dataloader.dataset import MTSDataset
-from utils.log import load_pkl
-from utils.log import TrainLogger
-from utils.load_data import re_max_min_normalization, standard_re_transform
-from losses.losses_torch import masked_mae, masked_rmse, masked_mape, metric
-
-
-class FullModelRunner(Runner):
- def __init__(self, cfg: dict, use_gpu: bool = True):
- super().__init__(cfg, use_gpu=use_gpu)
- self.clip = 3
- self._lambda = 1
- dataset_name = cfg['DATASET_NAME']
- # scaler
- if dataset_name in ['PEMS04', 'PEMS08']:
- _min = load_pkl("datasets/" + dataset_name + "/min.pkl")
- _max = load_pkl("datasets/" + dataset_name + "/max.pkl")
- self.scaler = re_max_min_normalization
- self.scaler_args = {'min': _min, 'max':_max}
- elif dataset_name in ['PEMS-BAY', 'METR-LA']:
- mean = load_pkl("datasets/" + dataset_name + "/mean.pkl")
- std = load_pkl("datasets/" + dataset_name + "/std.pkl")
- self.scaler = standard_re_transform
- self.scaler_args = {'mean': mean, 'std':std}
- self.loss = masked_mae
-
- self.dataset_name = cfg['DATASET_NAME']
- self.output_seq_len = 12
- self.cl_len = self.output_seq_len
- self.if_cl = True
-
- def init_training(self, cfg):
- """Initialize training.
-
- Including loss, training meters, etc.
-
- Args:
- cfg (dict): config
- """
- self.register_epoch_meter("train_loss", 'train', '{:.4f}')
- self.register_epoch_meter("train_MAPE", 'train', '{:.4f}')
- self.register_epoch_meter("train_RMSE", 'train', '{:.4f}')
-
- super().init_training(cfg)
-
- def init_validation(self, cfg: dict):
- """Initialize validation.
-
- Including validation meters, etc.
-
- Args:
- cfg (dict): config
- """
-
- super().init_validation(cfg)
-
- self.register_epoch_meter("val_loss", 'val', '{:.4f}')
- self.register_epoch_meter("val_MAPE", 'val', '{:.4f}')
- self.register_epoch_meter("val_RMSE", 'val', '{:.4f}')
-
- def init_test(self, cfg: dict):
- """Initialize test.
-
- Including test meters, etc.
-
- Args:
- cfg (dict): config
- """
-
- super().init_test(cfg)
-
- self.register_epoch_meter("test_loss", 'test', '{:.4f}')
- self.register_epoch_meter("test_MAPE", 'test', '{:.4f}')
- self.register_epoch_meter("test_RMSE", 'test', '{:.4f}')
-
- @staticmethod
- def define_model(cfg: dict) -> nn.Module:
- """Define model.
-
- If you have multiple models, insert the name and class into the dict below,
- and select it through ```config```.
-
- Args:
- cfg (dict): config
-
- Returns:
- model (nn.Module)
- """
- return {
- 'STEP': STEP
- }[cfg['MODEL']['NAME']](cfg, **cfg.MODEL.PARAM)
-
- def build_train_dataset(self, cfg: dict):
- """Build train dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["TRAIN"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["TRAIN"]["DATA"]["DIR"] + "/train_index.pkl"
- seq_len = cfg['TRAIN']['DATA']['SEQ_LEN']
- batch_size = cfg['TRAIN']['DATA']['BATCH_SIZE']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, throw=False, pretrain=False)
-
- warmup_epoches = cfg['TRAIN']['WARMUP_EPOCHS']
- cl_epochs = cfg['TRAIN']['CL_EPOCHS']
- self.init_lr = cfg['TRAIN']['OPTIM']['PARAM']['lr']
- self.itera_per_epoch = math.ceil(len(dataset) / batch_size)
- self.warmup_steps = self.itera_per_epoch * warmup_epoches
- self.cl_steps = self.itera_per_epoch * cl_epochs
- print("cl_steps:{0}".format(self.cl_steps))
- print("warmup_steps:{0}".format(self.warmup_steps))
-
- return dataset
-
- @staticmethod
- def build_val_dataset(cfg: dict):
- """Build val dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["VAL"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["VAL"]["DATA"]["DIR"] + "/valid_index.pkl"
- seq_len = cfg['VAL']['DATA']['SEQ_LEN']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, pretrain=False)
- print("val len: {0}".format(len(dataset)))
- return dataset
-
- @staticmethod
- def build_test_dataset(cfg: dict):
- """Build val dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["TEST"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["TEST"]["DATA"]["DIR"] + "/test_index.pkl"
- seq_len = cfg['TEST']['DATA']['SEQ_LEN']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, pretrain=False)
- print("test len: {0}".format(len(dataset)))
- return dataset
-
- def train_iters(self, epoch, iter_index, data):
- """Training details.
-
- Args:
- epoch (int): current epoch.
- iter_index (int): current iter.
- data (torch.Tensor or tuple): Data provided by DataLoader
-
- Returns:
- loss (torch.Tensor)
- """
- iter_num = (epoch-1) * self.itera_per_epoch + iter_index
-
- y, short_x, long_x = data
- Y = self.to_running_device(y)
- short_term_history = self.to_running_device(short_x)
- long_term_history = self.to_running_device(long_x)
-
- output, theta, priori_adj = self.model(short_term_history, long_term_history=long_term_history)
- output = output.transpose(1,2)
-
- # # reg
- B, N, N = theta.shape
- theta = theta.view(B, N*N)
- tru = priori_adj.view(B, N*N)
- BCE_loss = nn.BCELoss()
- loss_g = BCE_loss(theta, tru)
-
- # curriculum learning
- if iter_num < self.warmup_steps: # warm upping
- self.cl_len = self.output_seq_len
- elif iter_num == self.warmup_steps:
- # init curriculum learning
- self.cl_len = 1
- for param_group in self.optim.param_groups:
- param_group["lr"] = self.init_lr
- print("======== Start curriculum learning... reset the learning rate to {0}. ========".format(self.init_lr))
- else:
- # begin curriculum learning
- if (iter_num - self.warmup_steps) % self.cl_steps == 0 and self.cl_len <= self.output_seq_len:
- self.cl_len += int(self.if_cl)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- predict = self.scaler(output.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = self.scaler(Y.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = real_val[..., 0]
- mae_loss = self.loss(predict[:, :self.cl_len, :], real_val[:, :self.cl_len, :])
- else:
- ## inverse transform for both predict and real value.
- predict = self.scaler(output, **self.scaler_args)
- real_val = self.scaler(Y[:,:,:,0], **self.scaler_args)
- mae_loss = self.loss(predict[:, :self.cl_len, :], real_val[:, :self.cl_len, :], 0)
-
- if iter_num == 50:
- self._lambda = 1
- if iter_num == 100:
- self._lambda == 0.1
- if iter_num == 1000:
- self._lambda = 0.01
- self._lambda = 1 / (int(epoch/6)+1)
-
- loss = mae_loss + self._lambda * loss_g
- # metrics
- mape = masked_mape(predict,real_val,0.0)
- rmse = masked_rmse(predict,real_val,0.0)
-
- self.update_epoch_meter('train_loss', loss.item())
- self.update_epoch_meter('train_MAPE', mape.item())
- self.update_epoch_meter('train_RMSE', rmse.item())
-
- return loss
-
- def val_iters(self, iter_index, data):
- """Validation details.
-
- Args:
- iter_index (int): current iter.
- data (torch.Tensor or tuple): Data provided by DataLoader
- """
- y, short_x, long_x = data
- Y = self.to_running_device(y)
- short_term_history = self.to_running_device(short_x)
- long_term_history = self.to_running_device(long_x)
-
- output, _, _ = self.model(short_term_history, long_term_history=long_term_history)
- output = output.transpose(1,2)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- predict = self.scaler(output.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = self.scaler(Y.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = real_val[..., 0]
- mae_loss = self.loss(predict[:, :self.cl_len, :], real_val[:, :self.cl_len, :])
- else:
- ## inverse transform for both predict and real value.
- predict = self.scaler(output, **self.scaler_args)
- real_val = self.scaler(Y[:,:,:,0], **self.scaler_args)
-
- # metrics
- loss = self.loss(predict, real_val, 0.0)
- mape = masked_mape(predict,real_val,0.0)
- rmse = masked_rmse(predict,real_val,0.0)
-
- self.update_epoch_meter('val_loss', loss.item())
- self.update_epoch_meter('val_MAPE', mape.item())
- self.update_epoch_meter('val_RMSE', rmse.item())
-
- @torch.no_grad()
- @master_only
- def test(self, cfg: dict = None, train_epoch: int = None):
- """test model.
-
- Args:
- cfg (dict, optional): config
- train_epoch (int, optional): current epoch if in training process.
- """
-
- # init test if not in training process
- if train_epoch is None:
- self.init_test(cfg)
-
- self.on_test_start()
-
- test_start_time = time.time()
- self.model.eval()
-
- # test loop
- outputs = []
- y_list = []
- for iter_index, data in enumerate(self.test_data_loader):
- preds, testy = self.test_iters(iter_index, data)
- outputs.append(preds)
- y_list.append(testy)
- yhat = torch.cat(outputs,dim=0)
- y_list = torch.cat(y_list, dim=0)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- real_val = self.scaler(y_list.squeeze(-1), **self.scaler_args).transpose(1, 2)
- predict = self.scaler(yhat.unsqueeze(-1), **self.scaler_args).transpose(1, 2)
- real_val = real_val[..., 0]
- predict = predict[..., 0]
- else:
- ## inverse transform for both predict and real value.
- real_val = self.scaler(y_list[:,:,:,0], **self.scaler_args).transpose(1, 2)
- predict = self.scaler(yhat, **self.scaler_args).transpose(1, 2)
-
- # summarize the results.
- amae = []
- amape = []
- armse = []
-
- for i in range(12):
- # For horizon i, only calculate the metrics **at that time** slice here.
- pred = predict[:,:,i]
- real = real_val[:,:,i]
- dataset_name = self.dataset_name
- if dataset_name == 'PEMS04' or dataset_name == 'PEMS08': # traffic flow dataset follows mae metric used in ASTGNN.
- mae = mean_absolute_error(pred.cpu().numpy(), real.cpu().numpy())
- rmse = masked_rmse(pred, real, 0.0).item()
- mape = masked_mape(pred, real, 0.0).item()
- log = 'Evaluate best model on test data for horizon {:d}, Test MAE: {:.4f}, Test RMSE: {:.4f}, Test MAPE: {:.4f}'
- log = log.format(i+1, mae, rmse, mape)
- # print(log)
- amae.append(mae)
- amape.append(mape)
- armse.append(rmse)
- else: # traffic speed datasets follow the metrics released by GWNet and DCRNN.
- metrics = metric(pred,real)
- log = 'Evaluate best model on test data for horizon {:d}, Test MAE: {:.4f}, Test RMSE: {:.4f}, Test MAPE: {:.4f}'
- log = log.format(i+1, metrics[0], metrics[2], metrics[1])
- # print(log)
- amae.append(metrics[0]) # mae
- amape.append(metrics[1]) # mape
- armse.append(metrics[2]) # rmse
- self.logger.info(log)
-
- # *** TODO: not test yet ***
- import numpy as np
- log = 'On average over 12 horizons, Test MAE: {:.4f}, Test MAPE: {:.4f}, Test RMSE: {:.4f}'
- self.logger.info(log.format(np.mean(amae),np.mean(amape),np.mean(armse)))
- # **************************
-
- test_end_time = time.time()
- self.update_epoch_meter('test_time', test_start_time - test_end_time)
- # print val meters
- self.print_epoch_meters('test')
- if train_epoch is not None:
- # tensorboard plt meters
- self.plt_epoch_meters('test', train_epoch // self.test_interval)
-
- self.on_test_end()
-
- def test_iters(self, iter_index: int, data: torch.Tensor or tuple):
- y, short_x, long_x = data
- Y = self.to_running_device(y)
- short_term_history = self.to_running_device(short_x)
- long_term_history = self.to_running_device(long_x)
-
- output, _, _ = self.model(short_term_history, long_term_history=long_term_history)
- output = output.transpose(1,2)
- return output, Y
-
\ No newline at end of file
diff --git a/Runner_TSFormer.py b/Runner_TSFormer.py
deleted file mode 100644
index b04c0fa..0000000
--- a/Runner_TSFormer.py
+++ /dev/null
@@ -1,237 +0,0 @@
-from torch import nn
-
-from easytorch import Runner
-
-from dataloader.dataset import MTSDataset
-from models.tsformer.model import TSFormer
-from utils.log import TrainLogger, load_pkl
-from utils.load_data import re_max_min_normalization, standard_re_transform
-from losses.losses_torch import masked_mae, masked_rmse, masked_mape
-
-
-class TSFormerRunner(Runner):
- def __init__(self, cfg: dict, use_gpu: bool = True):
- super().__init__(cfg, use_gpu=use_gpu)
- self.clip = 5
- dataset_name = cfg['DATASET_NAME']
- # scaler
- if dataset_name in ['PEMS04', 'PEMS08']:
- _min = load_pkl("datasets/" + dataset_name + "/min.pkl")
- _max = load_pkl("datasets/" + dataset_name + "/max.pkl")
- self.scaler = re_max_min_normalization
- self.scaler_args = {'min': _min, 'max':_max}
- elif dataset_name in ['PEMS-BAY', 'METR-LA']:
- mean = load_pkl("datasets/" + dataset_name + "/mean.pkl")
- std = load_pkl("datasets/" + dataset_name + "/std.pkl")
- self.scaler = standard_re_transform
- self.scaler_args = {'mean': mean, 'std':std}
- self.loss = masked_mae
-
- def init_training(self, cfg):
- """Initialize training.
-
- Including loss, training meters, etc.
-
- Args:
- cfg (dict): config
- """
- super().init_training(cfg)
-
- self.register_epoch_meter("train_loss", 'train', '{:.4f}')
- self.register_epoch_meter("train_MAPE", 'train', '{:.4f}')
- self.register_epoch_meter("train_RMSE", 'train', '{:.4f}')
-
- def init_validation(self, cfg: dict):
- """Initialize validation.
-
- Including validation meters, etc.
-
- Args:
- cfg (dict): config
- """
-
- super().init_validation(cfg)
-
- self.register_epoch_meter("val_loss", 'val', '{:.4f}')
- self.register_epoch_meter("val_MAPE", 'val', '{:.4f}')
- self.register_epoch_meter("val_RMSE", 'val', '{:.4f}')
-
- def init_test(self, cfg: dict):
- """Initialize test.
-
- Including test meters, etc.
-
- Args:
- cfg (dict): config
- """
-
- super().init_test(cfg)
-
- self.register_epoch_meter("test_loss", 'test', '{:.4f}')
- self.register_epoch_meter("test_MAPE", 'test', '{:.4f}')
- self.register_epoch_meter("test_RMSE", 'test', '{:.4f}')
-
- @staticmethod
- def define_model(cfg: dict) -> nn.Module:
- """Define model.
-
- If you have multiple models, insert the name and class into the dict below,
- and select it through ```config```.
-
- Args:
- cfg (dict): config
-
- Returns:
- model (nn.Module)
- """
-
- return {
- 'TSFormer': TSFormer
- }[cfg['MODEL']['NAME']](**cfg['MODEL'].get('PARAM', {}))
-
- @staticmethod
- def build_train_dataset(cfg: dict):
- """Build MNIST train dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["TRAIN"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["TRAIN"]["DATA"]["DIR"] + "/train_index.pkl"
- seq_len = cfg['TRAIN']['DATA']['SEQ_LEN']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, throw=True, pretrain=True)
- return dataset
-
- @staticmethod
- def build_val_dataset(cfg: dict):
- """Build MNIST val dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["VAL"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["VAL"]["DATA"]["DIR"] + "/valid_index.pkl"
- seq_len = cfg['VAL']['DATA']['SEQ_LEN']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, throw=True, pretrain=True)
- return dataset
-
- @staticmethod
- def build_test_dataset(cfg: dict):
- """Build MNIST val dataset
-
- Args:
- cfg (dict): config
-
- Returns:
- train dataset (Dataset)
- """
- raw_file_path = cfg["TEST"]["DATA"]["DIR"] + "/data.pkl"
- index_file_path = cfg["TEST"]["DATA"]["DIR"] + "/test_index.pkl"
- seq_len = cfg['TEST']['DATA']['SEQ_LEN']
- dataset = MTSDataset(raw_file_path, index_file_path, seq_len, throw=True, pretrain=True)
- return dataset
-
- def train_iters(self, epoch, iter_index, data):
- """Training details.
-
- Args:
- epoch (int): current epoch.
- iter_index (int): current iter.
- data (torch.Tensor or tuple): Data provided by DataLoader
-
- Returns:
- loss (torch.Tensor)
- """
- input_data = self.to_running_device(data[0])
- # abs_idx = self.to_running_device(data[1]).long()
-
- output_masked_tokens, label_masked_tokens, plot_args = self.model(input_data)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- predict = self.scaler(output_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = self.scaler(label_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- mae_loss = self.loss(predict, real_val)
- else:
- ## inverse transform for both predict and real value.
- predict = self.scaler(output_masked_tokens, **self.scaler_args)
- real_val = self.scaler(label_masked_tokens, **self.scaler_args)
- mae_loss = self.loss(predict, real_val, 0)
-
- loss = mae_loss
- # metrics
- mape = masked_mape(predict,real_val,0.0)
- rmse = masked_rmse(predict,real_val,0.0)
-
- self.update_epoch_meter('train_loss', loss.item())
- self.update_epoch_meter('train_MAPE', mape.item())
- self.update_epoch_meter('train_RMSE', rmse.item())
- return loss
-
- def val_iters(self, iter_index, data):
- """Validation details.
-
- Args:
- iter_index (int): current iter.
- data (torch.Tensor or tuple): Data provided by DataLoader
- """
- input_data = self.to_running_device(data[0])
- # abs_idx = self.to_running_device(data[1]).long()
-
- output_masked_tokens, label_masked_tokens, plot_args = self.model(input_data)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- predict = self.scaler(output_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = self.scaler(label_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- loss = self.loss(predict, real_val)
- else:
- ## inverse transform for both predict and real value.
- predict = self.scaler(output_masked_tokens, **self.scaler_args)
- real_val = self.scaler(label_masked_tokens, **self.scaler_args)
- loss = self.loss(predict, real_val, 0.0)
-
- # metrics
- mape = masked_mape(predict,real_val,0.0)
- rmse = masked_rmse(predict,real_val,0.0)
-
- self.update_epoch_meter('val_loss', loss.item())
- self.update_epoch_meter('val_MAPE', mape.item())
- self.update_epoch_meter('val_RMSE', rmse.item())
-
- def test_iters(self, iter_index, data):
- """Validation details.
-
- Args:
- iter_index (int): current iter.
- data (torch.Tensor or tuple): Data provided by DataLoader
- """
- input_data = self.to_running_device(data[0])
- # abs_idx = self.to_running_device(data[1]).long()
-
- output_masked_tokens, label_masked_tokens, plot_args = self.model(input_data)
-
- # scale data and calculate loss
- if "max" in self.scaler_args.keys(): # traffic flow
- predict = self.scaler(output_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- real_val = self.scaler(label_masked_tokens.transpose(1,2).unsqueeze(-1), **self.scaler_args).transpose(1, 2).squeeze(-1)
- loss = self.loss(predict, real_val)
- else:
- ## inverse transform for both predict and real value.
- predict = self.scaler(output_masked_tokens, **self.scaler_args)
- real_val = self.scaler(label_masked_tokens, **self.scaler_args)
- loss = self.loss(predict, real_val, 0.0)
-
- # metrics
- mape = masked_mape(predict,real_val,0.0)
- rmse = masked_rmse(predict,real_val,0.0)
-
- self.update_epoch_meter('test_loss', loss.item())
- self.update_epoch_meter('test_MAPE', mape.item())
- self.update_epoch_meter('test_RMSE', rmse.item())
diff --git a/TODO.md b/TODO.md
new file mode 100644
index 0000000..f95d115
--- /dev/null
+++ b/TODO.md
@@ -0,0 +1 @@
+- Update README.md and datasets/README.md
\ No newline at end of file
diff --git a/basicts/__init__.py b/basicts/__init__.py
new file mode 100644
index 0000000..7cc0167
--- /dev/null
+++ b/basicts/__init__.py
@@ -0,0 +1,5 @@
+from .launcher import launch_training
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__", "launch_training"]
diff --git a/basicts/archs/__init__.py b/basicts/archs/__init__.py
new file mode 100644
index 0000000..a06aab3
--- /dev/null
+++ b/basicts/archs/__init__.py
@@ -0,0 +1,25 @@
+from .arch_zoo.stid_arch import STID
+from .arch_zoo.gwnet_arch import GraphWaveNet
+from .arch_zoo.dcrnn_arch import DCRNN
+from .arch_zoo.d2stgnn_arch import D2STGNN
+from .arch_zoo.stgcn_arch import STGCN
+from .arch_zoo.mtgnn_arch import MTGNN
+from .arch_zoo.stnorm_arch import STNorm
+from .arch_zoo.agcrn_arch import AGCRN
+from .arch_zoo.stemgnn_arch import StemGNN
+from .arch_zoo.gts_arch import GTS
+from .arch_zoo.dgcrn_arch import DGCRN
+from .arch_zoo.linear_arch import Linear, DLinear, NLinear
+from .arch_zoo.autoformer_arch import Autoformer
+from .arch_zoo.hi_arch import HINetwork
+from .arch_zoo.fedformer_arch import FEDformer
+from .arch_zoo.informer_arch import Informer, InformerStack
+from .arch_zoo.pyraformer_arch import Pyraformer
+
+__all__ = ["STID", "GraphWaveNet", "DCRNN",
+ "D2STGNN", "STGCN", "MTGNN",
+ "STNorm", "AGCRN", "StemGNN",
+ "GTS", "DGCRN", "Linear",
+ "DLinear", "NLinear", "Autoformer",
+ "HINetwork", "FEDformer", "Informer",
+ "InformerStack", "Pyraformer"]
diff --git a/basicts/archs/arch_zoo/agcrn_arch/__init__.py b/basicts/archs/arch_zoo/agcrn_arch/__init__.py
new file mode 100644
index 0000000..01b4cc8
--- /dev/null
+++ b/basicts/archs/arch_zoo/agcrn_arch/__init__.py
@@ -0,0 +1,3 @@
+from .agcrn_arch import AGCRN
+
+__all__ = ["AGCRN"]
diff --git a/basicts/archs/arch_zoo/agcrn_arch/agcn.py b/basicts/archs/arch_zoo/agcrn_arch/agcn.py
new file mode 100644
index 0000000..1cea52e
--- /dev/null
+++ b/basicts/archs/arch_zoo/agcrn_arch/agcn.py
@@ -0,0 +1,35 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+
+class AVWGCN(nn.Module):
+ def __init__(self, dim_in, dim_out, cheb_k, embed_dim):
+ super(AVWGCN, self).__init__()
+ self.cheb_k = cheb_k
+ self.weights_pool = nn.Parameter(
+ torch.FloatTensor(embed_dim, cheb_k, dim_in, dim_out))
+ self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, dim_out))
+
+ def forward(self, x, node_embeddings):
+ # x shaped[B, N, C], node_embeddings shaped [N, D] -> supports shaped [N, N]
+ # output shape [B, N, C]
+ node_num = node_embeddings.shape[0]
+ supports = F.softmax(
+ F.relu(torch.mm(node_embeddings, node_embeddings.transpose(0, 1))), dim=1)
+ support_set = [torch.eye(node_num).to(supports.device), supports]
+ # default cheb_k = 3
+ for k in range(2, self.cheb_k):
+ support_set.append(torch.matmul(
+ 2 * supports, support_set[-1]) - support_set[-2])
+ supports = torch.stack(support_set, dim=0)
+ # N, cheb_k, dim_in, dim_out
+ weights = torch.einsum(
+ 'nd,dkio->nkio', node_embeddings, self.weights_pool)
+ bias = torch.matmul(node_embeddings, self.bias_pool) # N, dim_out
+ x_g = torch.einsum("knm,bmc->bknc", supports,
+ x) # B, cheb_k, N, dim_in
+ x_g = x_g.permute(0, 2, 1, 3) # B, N, cheb_k, dim_in
+ x_gconv = torch.einsum('bnki,nkio->bno', x_g,
+ weights) + bias # b, N, dim_out
+ return x_gconv
diff --git a/basicts/archs/arch_zoo/agcrn_arch/agcrn_arch.py b/basicts/archs/arch_zoo/agcrn_arch/agcrn_arch.py
new file mode 100644
index 0000000..a64a1c8
--- /dev/null
+++ b/basicts/archs/arch_zoo/agcrn_arch/agcrn_arch.py
@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+
+from .agcrn_cell import AGCRNCell
+
+
+class AVWDCRNN(nn.Module):
+ def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim, num_layers=1):
+ super(AVWDCRNN, self).__init__()
+ assert num_layers >= 1, 'At least one DCRNN layer in the Encoder.'
+ self.node_num = node_num
+ self.input_dim = dim_in
+ self.num_layers = num_layers
+ self.dcrnn_cells = nn.ModuleList()
+ self.dcrnn_cells.append(
+ AGCRNCell(node_num, dim_in, dim_out, cheb_k, embed_dim))
+ for _ in range(1, num_layers):
+ self.dcrnn_cells.append(
+ AGCRNCell(node_num, dim_out, dim_out, cheb_k, embed_dim))
+
+ def forward(self, x, init_state, node_embeddings):
+ # shape of x: (B, T, N, D)
+ # shape of init_state: (num_layers, B, N, hidden_dim)
+ assert x.shape[2] == self.node_num and x.shape[3] == self.input_dim
+ seq_length = x.shape[1]
+ current_inputs = x
+ output_hidden = []
+ for i in range(self.num_layers):
+ state = init_state[i]
+ inner_states = []
+ for t in range(seq_length):
+ state = self.dcrnn_cells[i](
+ current_inputs[:, t, :, :], state, node_embeddings)
+ inner_states.append(state)
+ output_hidden.append(state)
+ current_inputs = torch.stack(inner_states, dim=1)
+ # current_inputs: the outputs of last layer: (B, T, N, hidden_dim)
+ # output_hidden: the last state for each layer: (num_layers, B, N, hidden_dim)
+ #last_state: (B, N, hidden_dim)
+ return current_inputs, output_hidden
+
+ def init_hidden(self, batch_size):
+ init_states = []
+ for i in range(self.num_layers):
+ init_states.append(
+ self.dcrnn_cells[i].init_hidden_state(batch_size))
+ # (num_layers, B, N, hidden_dim)
+ return torch.stack(init_states, dim=0)
+
+
+class AGCRN(nn.Module):
+ """
+ Paper: Adaptive Graph Convolutional Recurrent Network for Traf๏ฌc Forecasting
+ Official Code: https://github.com/LeiBAI/AGCRN
+ Link: https://arxiv.org/abs/2007.02842
+ """
+
+ def __init__(self, num_nodes, input_dim, rnn_units, output_dim, horizon, num_layers, default_graph, embed_dim, cheb_k):
+ super(AGCRN, self).__init__()
+ self.num_node = num_nodes
+ self.input_dim = input_dim
+ self.hidden_dim = rnn_units
+ self.output_dim = output_dim
+ self.horizon = horizon
+ self.num_layers = num_layers
+
+ self.default_graph = default_graph
+ self.node_embeddings = nn.Parameter(torch.randn(
+ self.num_node, embed_dim), requires_grad=True)
+
+ self.encoder = AVWDCRNN(num_nodes, input_dim, rnn_units, cheb_k,
+ embed_dim, num_layers)
+
+ # predictor
+ self.end_conv = nn.Conv2d(
+ 1, horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
+
+ self.init_param()
+
+ def init_param(self):
+ for p in self.parameters():
+ if p.dim() > 1:
+ nn.init.xavier_uniform_(p)
+ else:
+ nn.init.uniform_(p)
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool, **kwargs) -> torch.Tensor:
+ """Feedforward function of AGCRN.
+
+ Args:
+ history_data (torch.Tensor): inputs with shape [B, L, N, C].
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L, N, C]
+ """
+
+ init_state = self.encoder.init_hidden(history_data.shape[0])
+ output, _ = self.encoder(
+ history_data, init_state, self.node_embeddings) # B, T, N, hidden
+ output = output[:, -1:, :, :] # B, 1, N, hidden
+
+ # CNN based predictor
+ output = self.end_conv((output)) # B, T*C, N, 1
+ output = output.squeeze(-1).reshape(-1, self.horizon,
+ self.output_dim, self.num_node)
+ output = output.permute(0, 1, 3, 2) # B, T, N, C
+
+ return output
diff --git a/basicts/archs/arch_zoo/agcrn_arch/agcrn_cell.py b/basicts/archs/arch_zoo/agcrn_arch/agcrn_cell.py
new file mode 100644
index 0000000..d11c6d5
--- /dev/null
+++ b/basicts/archs/arch_zoo/agcrn_arch/agcrn_cell.py
@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+
+from .agcn import AVWGCN
+
+
+class AGCRNCell(nn.Module):
+ def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim):
+ super(AGCRNCell, self).__init__()
+ self.node_num = node_num
+ self.hidden_dim = dim_out
+ self.gate = AVWGCN(dim_in+self.hidden_dim, 2 *
+ dim_out, cheb_k, embed_dim)
+ self.update = AVWGCN(dim_in+self.hidden_dim,
+ dim_out, cheb_k, embed_dim)
+
+ def forward(self, x, state, node_embeddings):
+ # x: B, num_nodes, input_dim
+ # state: B, num_nodes, hidden_dim
+ state = state.to(x.device)
+ input_and_state = torch.cat((x, state), dim=-1)
+ z_r = torch.sigmoid(self.gate(input_and_state, node_embeddings))
+ z, r = torch.split(z_r, self.hidden_dim, dim=-1)
+ candidate = torch.cat((x, z*state), dim=-1)
+ hc = torch.tanh(self.update(candidate, node_embeddings))
+ h = r*state + (1-r)*hc
+ return h
+
+ def init_hidden_state(self, batch_size):
+ return torch.zeros(batch_size, self.node_num, self.hidden_dim)
diff --git a/basicts/archs/arch_zoo/autoformer_arch/__init__.py b/basicts/archs/arch_zoo/autoformer_arch/__init__.py
new file mode 100644
index 0000000..d9d867b
--- /dev/null
+++ b/basicts/archs/arch_zoo/autoformer_arch/__init__.py
@@ -0,0 +1,3 @@
+from .autoformer_arch import Autoformer
+
+__all__ = ["Autoformer"]
diff --git a/basicts/archs/arch_zoo/autoformer_arch/auto_correlation.py b/basicts/archs/arch_zoo/autoformer_arch/auto_correlation.py
new file mode 100644
index 0000000..d62f93a
--- /dev/null
+++ b/basicts/archs/arch_zoo/autoformer_arch/auto_correlation.py
@@ -0,0 +1,160 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+class AutoCorrelation(nn.Module):
+ """
+ AutoCorrelation Mechanism with the following two phases:
+ (1) period-based dependencies discovery
+ (2) time delay aggregation
+ This block can replace the self-attention family mechanism seamlessly.
+ """
+ def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False):
+ super(AutoCorrelation, self).__init__()
+ self.factor = factor
+ self.scale = scale
+ self.mask_flag = mask_flag
+ self.output_attention = output_attention
+ self.dropout = nn.Dropout(attention_dropout)
+
+ def time_delay_agg_training(self, values, corr):
+ """
+ SpeedUp version of Autocorrelation (a batch-normalization style design)
+ This is for the training phase.
+ """
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+ index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
+ weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ pattern = torch.roll(tmp_values, -int(index[i]), -1)
+ delays_agg = delays_agg + pattern * \
+ (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
+ return delays_agg
+
+ def time_delay_agg_inference(self, values, corr):
+ """
+ SpeedUp version of Autocorrelation (a batch-normalization style design)
+ This is for the inference phase.
+ """
+ batch = values.shape[0]
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # index init
+ init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0)\
+ .repeat(batch, head, channel, 1).to(values.device)
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+ weights, delay = torch.topk(mean_value, top_k, dim=-1)
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values.repeat(1, 1, 1, 2)
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
+ pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+ delays_agg = delays_agg + pattern * \
+ (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
+ return delays_agg
+
+ def time_delay_agg_full(self, values, corr):
+ """
+ Standard version of Autocorrelation
+ """
+ batch = values.shape[0]
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # index init
+ init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0)\
+ .repeat(batch, head, channel, 1).to(values.device)
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ weights, delay = torch.topk(corr, top_k, dim=-1)
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values.repeat(1, 1, 1, 2)
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ tmp_delay = init_index + delay[..., i].unsqueeze(-1)
+ pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+ delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
+ return delays_agg
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, H, E = queries.shape
+ _, S, _, D = values.shape
+ if L > S:
+ zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
+ values = torch.cat([values, zeros], dim=1)
+ keys = torch.cat([keys, zeros], dim=1)
+ else:
+ values = values[:, :L, :, :]
+ keys = keys[:, :L, :, :]
+
+ # period-based dependencies
+ q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1)
+ k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
+ res = q_fft * torch.conj(k_fft)
+ corr = torch.fft.irfft(res, dim=-1)
+
+ # time delay agg
+ if self.training:
+ V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+ else:
+ V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+
+ if self.output_attention:
+ return (V.contiguous(), corr.permute(0, 3, 1, 2))
+ else:
+ return (V.contiguous(), None)
+
+
+class AutoCorrelationLayer(nn.Module):
+ def __init__(self, correlation, d_model, n_heads, d_keys=None,
+ d_values=None):
+ super(AutoCorrelationLayer, self).__init__()
+
+ d_keys = d_keys or (d_model // n_heads)
+ d_values = d_values or (d_model // n_heads)
+
+ self.inner_correlation = correlation
+ self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.value_projection = nn.Linear(d_model, d_values * n_heads)
+ self.out_projection = nn.Linear(d_values * n_heads, d_model)
+ self.n_heads = n_heads
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, _ = queries.shape
+ _, S, _ = keys.shape
+ H = self.n_heads
+
+ queries = self.query_projection(queries).view(B, L, H, -1)
+ keys = self.key_projection(keys).view(B, S, H, -1)
+ values = self.value_projection(values).view(B, S, H, -1)
+
+ out, attn = self.inner_correlation(
+ queries,
+ keys,
+ values,
+ attn_mask
+ )
+ out = out.view(B, L, -1)
+
+ return self.out_projection(out), attn
diff --git a/basicts/archs/arch_zoo/autoformer_arch/autoformer_arch.py b/basicts/archs/arch_zoo/autoformer_arch/autoformer_arch.py
new file mode 100644
index 0000000..ec3e249
--- /dev/null
+++ b/basicts/archs/arch_zoo/autoformer_arch/autoformer_arch.py
@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+
+from .embed import DataEmbedding_wo_pos, DataEmbedding
+from .auto_correlation import AutoCorrelation, AutoCorrelationLayer
+from .enc_dec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp
+from ..utils import data_transformation_4_xformer
+
+
+class Autoformer(nn.Module):
+ """
+ Paper: Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting
+ Link: https://arxiv.org/abs/2206.09112
+ Ref Official Code: https://github.com/thuml/Autoformer
+ """
+
+ def __init__(self, **model_args):
+ super(Autoformer, self).__init__()
+ self.seq_len = int(model_args["seq_len"])
+ self.label_len = int(model_args["label_len"])
+ self.pred_len = int(model_args["pred_len"])
+ self.output_attention = model_args["output_attention"]
+ self.embedding_type = model_args["embedding_type"]
+
+ # Decomp
+ kernel_size = model_args["moving_avg"]
+ self.decomp = series_decomp(kernel_size)
+
+ # Embedding
+ # The series-wise connection inherently contains the sequential information.
+ # Thus, we can discard the position embedding of transformers.
+ if self.embedding_type == "DataEmbedding_wo_pos":
+ self.enc_embedding = DataEmbedding_wo_pos(model_args["enc_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ self.dec_embedding = DataEmbedding_wo_pos(model_args["dec_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ elif self.embedding_type == "DataEmbedding":
+ self.enc_embedding = DataEmbedding(model_args["enc_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ self.dec_embedding = DataEmbedding(model_args["dec_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ else:
+ raise Exception("Unknown embedding type.")
+
+ # Encoder
+ self.encoder = Encoder(
+ [
+ EncoderLayer(
+ AutoCorrelationLayer(
+ AutoCorrelation(False, model_args["factor"], attention_dropout=model_args["dropout"],
+ output_attention=model_args["output_attention"]),
+ model_args["d_model"], model_args["n_heads"]),
+ model_args["d_model"],
+ model_args["d_ff"],
+ moving_avg=model_args["moving_avg"],
+ dropout=model_args["dropout"],
+ activation=model_args["activation"]
+ ) for l in range(model_args["e_layers"])
+ ],
+ norm_layer=my_Layernorm(model_args["d_model"])
+ )
+ # Decoder
+ self.decoder = Decoder(
+ [
+ DecoderLayer(
+ AutoCorrelationLayer(
+ AutoCorrelation(True, model_args["factor"], attention_dropout=model_args["dropout"],
+ output_attention=False),
+ model_args["d_model"], model_args["n_heads"]),
+ AutoCorrelationLayer(
+ AutoCorrelation(False, model_args["factor"], attention_dropout=model_args["dropout"],
+ output_attention=False),
+ model_args["d_model"], model_args["n_heads"]),
+ model_args["d_model"],
+ model_args["c_out"],
+ model_args["d_ff"],
+ moving_avg=model_args["moving_avg"],
+ dropout=model_args["dropout"],
+ activation=model_args["activation"],
+ )
+ for l in range(model_args["d_layers"])
+ ],
+ norm_layer=my_Layernorm(model_args["d_model"]),
+ projection=nn.Linear(model_args["d_model"], model_args["c_out"], bias=True)
+ )
+
+ def forward_xformer(self, x_enc: torch.Tensor, x_mark_enc: torch.Tensor, x_dec: torch.Tensor, x_mark_dec: torch.Tensor,
+ enc_self_mask: torch.Tensor=None, dec_self_mask: torch.Tensor=None, dec_enc_mask: torch.Tensor=None) -> torch.Tensor:
+ """Feed forward of AutoFormer. Kindly note that `enc_self_mask`, `dec_self_mask`, and `dec_enc_mask` are not actually used in AutoFormer.
+ See:
+ - https://github.com/thuml/Autoformer/blob/e116bbcf41f537f4ab53d172d9babfc0a026330f/layers/AutoCorrelation.py#L103
+ - https://github.com/thuml/Autoformer/blob/e116bbcf41f537f4ab53d172d9babfc0a026330f/exp/exp_main.py#L136
+
+ Args:
+ x_enc (torch.Tensor): input data of encoder (without the time features). Shape: [B, L1, N]
+ x_mark_enc (torch.Tensor): time features input of encoder w.r.t. x_enc. Shape: [B, L1, C-1]
+ x_dec (torch.Tensor): input data of decoder. Shape: [B, start_token_length + L2, N]
+ x_mark_dec (torch.Tensor): time features input to decoder w.r.t. x_dec. Shape: [B, start_token_length + L2, C-1]
+ enc_self_mask (torch.Tensor, optional): encoder self attention masks. Defaults to None.
+ dec_self_mask (torch.Tensor, optional): decoder self attention masks. Defaults to None.
+ dec_enc_mask (torch.Tensor, optional): decoder encoder self attention masks. Defaults to None.
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L2, N, 1]
+ """
+
+ # decomp init
+ mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
+ zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)
+ seasonal_init, trend_init = self.decomp(x_enc)
+ # decoder input
+ trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
+ seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)
+ # enc
+ enc_out = self.enc_embedding(x_enc, x_mark_enc)
+ enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+ # dec
+ dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
+ seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
+ trend=trend_init)
+ # final
+ dec_out = trend_part + seasonal_part
+
+ return dec_out[:, -self.pred_len:, :].unsqueeze(-1) # [B, L, N, C]
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool, **kwargs) -> torch.Tensor:
+ """
+
+ Args:
+ history_data (Tensor): Input data with shape: [B, L1, N, C]
+ future_data (Tensor): Future data with shape: [B, L2, N, C]
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L2, N, 1]
+ """
+
+ x_enc, x_mark_enc, x_dec, x_mark_dec = data_transformation_4_xformer(history_data=history_data, future_data=future_data, start_token_len=self.label_len)
+ prediction = self.forward_xformer(x_enc=x_enc, x_mark_enc=x_mark_enc, x_dec=x_dec, x_mark_dec=x_mark_dec)
+ return prediction
diff --git a/basicts/archs/arch_zoo/autoformer_arch/embed.py b/basicts/archs/arch_zoo/autoformer_arch/embed.py
new file mode 100644
index 0000000..c5ccb01
--- /dev/null
+++ b/basicts/archs/arch_zoo/autoformer_arch/embed.py
@@ -0,0 +1,98 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+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, num_time_features):
+ super(TimeFeatureEmbedding, self).__init__()
+ # freq_map = {'h': 2}
+ # NOTE: this is different from official autoformer, since METR-LA, PEMS-BAY, and PEMS0X datasets only contain two kind of temporal features (time in day, day in week)
+ self.embed = nn.Linear(num_time_features, d_model, bias=False)
+
+ def forward(self, x):
+ return self.embed(x)
+
+
+class DataEmbedding(nn.Module):
+ def __init__(self, c_in, d_model, num_time_features, 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=d_model, num_time_features=num_time_features)
+ 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)
+
+
+class DataEmbedding_wo_pos(nn.Module):
+ def __init__(self, c_in, d_model, num_time_features, dropout=0.1):
+ super(DataEmbedding_wo_pos, 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=d_model, num_time_features=num_time_features)
+ self.dropout = nn.Dropout(p=dropout)
+
+ def forward(self, x, x_mark):
+ x = self.value_embedding(x) + self.temporal_embedding(x_mark)
+ return self.dropout(x)
diff --git a/basicts/archs/arch_zoo/autoformer_arch/enc_dec.py b/basicts/archs/arch_zoo/autoformer_arch/enc_dec.py
new file mode 100644
index 0000000..00b3a7a
--- /dev/null
+++ b/basicts/archs/arch_zoo/autoformer_arch/enc_dec.py
@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class my_Layernorm(nn.Module):
+ """
+ Special designed layernorm for the seasonal part
+ """
+ def __init__(self, channels):
+ super(my_Layernorm, self).__init__()
+ self.layernorm = nn.LayerNorm(channels)
+
+ def forward(self, x):
+ x_hat = self.layernorm(x)
+ bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
+ return x_hat - bias
+
+
+class moving_avg(nn.Module):
+ """
+ Moving average block to highlight the trend of time series
+ """
+ def __init__(self, kernel_size, stride):
+ super(moving_avg, self).__init__()
+ self.kernel_size = kernel_size
+ self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+ def forward(self, x):
+ # padding on the both ends of time series
+ front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+ end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+ x = torch.cat([front, x, end], dim=1)
+ x = self.avg(x.permute(0, 2, 1))
+ x = x.permute(0, 2, 1)
+ return x
+
+
+class series_decomp(nn.Module):
+ """
+ Series decomposition block
+ """
+ def __init__(self, kernel_size):
+ super(series_decomp, self).__init__()
+ self.moving_avg = moving_avg(kernel_size, stride=1)
+
+ def forward(self, x):
+ moving_mean = self.moving_avg(x)
+ res = x - moving_mean
+ return res, moving_mean
+
+
+class EncoderLayer(nn.Module):
+ """
+ Autoformer encoder layer with the progressive decomposition architecture
+ """
+ def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
+ super(EncoderLayer, self).__init__()
+ d_ff = d_ff or 4 * d_model
+ self.attention = attention
+ self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+ self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+ self.decomp1 = series_decomp(moving_avg)
+ self.decomp2 = series_decomp(moving_avg)
+ self.dropout = nn.Dropout(dropout)
+ self.activation = F.relu if activation == "relu" else F.gelu
+
+ def forward(self, x, attn_mask=None):
+ new_x, attn = self.attention(
+ x, x, x,
+ attn_mask=attn_mask
+ )
+ x = x + self.dropout(new_x)
+ x, _ = self.decomp1(x)
+ y = x
+ y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+ y = self.dropout(self.conv2(y).transpose(-1, 1))
+ res, _ = self.decomp2(x + y)
+ return res, attn
+
+
+class Encoder(nn.Module):
+ """
+ Autoformer encoder
+ """
+ def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+ super(Encoder, self).__init__()
+ self.attn_layers = nn.ModuleList(attn_layers)
+ self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+ self.norm = norm_layer
+
+ def forward(self, x, attn_mask=None):
+ attns = []
+ if self.conv_layers is not None:
+ for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
+ x, attn = attn_layer(x, attn_mask=attn_mask)
+ x = conv_layer(x)
+ attns.append(attn)
+ x, attn = self.attn_layers[-1](x)
+ attns.append(attn)
+ else:
+ for attn_layer in self.attn_layers:
+ x, attn = attn_layer(x, attn_mask=attn_mask)
+ attns.append(attn)
+
+ if self.norm is not None:
+ x = self.norm(x)
+
+ return x, attns
+
+
+class DecoderLayer(nn.Module):
+ """
+ Autoformer decoder layer with the progressive decomposition architecture
+ """
+ def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
+ moving_avg=25, dropout=0.1, activation="relu"):
+ super(DecoderLayer, self).__init__()
+ d_ff = d_ff or 4 * d_model
+ self.self_attention = self_attention
+ self.cross_attention = cross_attention
+ self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+ self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+ self.decomp1 = series_decomp(moving_avg)
+ self.decomp2 = series_decomp(moving_avg)
+ self.decomp3 = series_decomp(moving_avg)
+ self.dropout = nn.Dropout(dropout)
+ self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
+ padding_mode='circular', bias=False)
+ self.activation = F.relu if activation == "relu" else F.gelu
+
+ def forward(self, x, cross, x_mask=None, cross_mask=None):
+ x = x + self.dropout(self.self_attention(
+ x, x, x,
+ attn_mask=x_mask
+ )[0])
+ x, trend1 = self.decomp1(x)
+ x = x + self.dropout(self.cross_attention(
+ x, cross, cross,
+ attn_mask=cross_mask
+ )[0])
+ x, trend2 = self.decomp2(x)
+ y = x
+ y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+ y = self.dropout(self.conv2(y).transpose(-1, 1))
+ x, trend3 = self.decomp3(x + y)
+
+ residual_trend = trend1 + trend2 + trend3
+ residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
+ return x, residual_trend
+
+
+class Decoder(nn.Module):
+ """
+ Autoformer encoder
+ """
+ def __init__(self, layers, norm_layer=None, projection=None):
+ super(Decoder, self).__init__()
+ self.layers = nn.ModuleList(layers)
+ self.norm = norm_layer
+ self.projection = projection
+
+ def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
+ for layer in self.layers:
+ x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+ trend = trend + residual_trend
+
+ if self.norm is not None:
+ x = self.norm(x)
+
+ if self.projection is not None:
+ x = self.projection(x)
+ return x, trend
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/__init__.py b/basicts/archs/arch_zoo/d2stgnn_arch/__init__.py
new file mode 100644
index 0000000..50fbf7f
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/__init__.py
@@ -0,0 +1,3 @@
+from .d2stgnn_arch import D2STGNN
+
+__all__ = ["D2STGNN"]
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/d2stgnn_arch.py b/basicts/archs/arch_zoo/d2stgnn_arch/d2stgnn_arch.py
new file mode 100644
index 0000000..91a4633
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/d2stgnn_arch.py
@@ -0,0 +1,184 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .difusion_block import DifBlock
+from .inherent_block import InhBlock
+from .dynamic_graph_conv.dy_graph_conv import DynamicGraphConstructor
+from .decouple.estimation_gate import EstimationGate
+
+
+class DecoupleLayer(nn.Module):
+ def __init__(self, hidden_dim, fk_dim=256, first=False, **model_args):
+ super().__init__()
+ self.spatial_gate = EstimationGate(
+ model_args['node_hidden'], model_args['time_emb_dim'], 64, model_args['seq_length'])
+ self.dif_layer = DifBlock(hidden_dim, fk_dim=fk_dim, **model_args)
+ self.inh_layer = InhBlock(
+ hidden_dim, fk_dim=fk_dim, first=first, **model_args)
+
+ def forward(self, X: torch.Tensor, dynamic_graph: torch.Tensor, static_graph, E_u, E_d, T_D, D_W):
+ """decouple layer
+
+ Args:
+ X (torch.Tensor): input data with shape (B, L, N, D)
+ dynamic_graph (list of torch.Tensor): dynamic graph adjacency matrix with shape (B, N, k_t * N)
+ static_graph (ist of torch.Tensor): the self-adaptive transition matrix with shape (N, N)
+ E_u (torch.Parameter): node embedding E_u
+ E_d (torch.Parameter): node embedding E_d
+ T_D (torch.Parameter): time embedding T_D
+ D_W (torch.Parameter): time embedding D_W
+
+ Returns:
+ torch.Tensor: the undecoupled signal in this layer, i.e., the X^{l+1}, which should be feeded to the next layer. shape [B, L', N, D].
+ torch.Tensor: the output of the forecast branch of Diffusion Block with shape (B, L'', N, D), where L''=output_seq_len / model_args['gap'] to avoid error accumulation in auto-regression.
+ torch.Tensor: the output of the forecast branch of Inherent Block with shape (B, L'', N, D), where L''=output_seq_len / model_args['gap'] to avoid error accumulation in auto-regression.
+ """
+ X_spa = self.spatial_gate(E_u, E_d, T_D, D_W, X)
+ dif_backcast_seq_res, dif_forecast_hidden = self.dif_layer(
+ X=X, X_spa=X_spa, dynamic_graph=dynamic_graph, static_graph=static_graph)
+ inh_backcast_seq_res, inh_forecast_hidden = self.inh_layer(
+ dif_backcast_seq_res)
+ return inh_backcast_seq_res, dif_forecast_hidden, inh_forecast_hidden
+
+
+class D2STGNN(nn.Module):
+ """
+ Paper: Decoupled Dynamic Spatial-Temporal Graph Neural Network for Traffic Forecasting
+ Link: https://arxiv.org/abs/2206.09112
+ Official Code: https://github.com/zezhishao/D2STGNN
+ """
+ def __init__(self, **model_args):
+ super().__init__()
+ # attributes
+ self._in_feat = model_args['num_feat']
+ self._hidden_dim = model_args['num_hidden']
+ self._node_dim = model_args['node_hidden']
+ self._forecast_dim = 256
+ self._output_hidden = 512
+ self._output_dim = model_args['seq_length']
+
+ self._num_nodes = model_args['num_nodes']
+ self._k_s = model_args['k_s']
+ self._k_t = model_args['k_t']
+ self._num_layers = 5
+
+ model_args['use_pre'] = False
+ model_args['dy_graph'] = True
+ model_args['sta_graph'] = True
+
+ self._model_args = model_args
+
+ # start embedding layer
+ self.embedding = nn.Linear(self._in_feat, self._hidden_dim)
+
+ # time embedding
+ self.T_i_D_emb = nn.Parameter(
+ torch.empty(288, model_args['time_emb_dim']))
+ self.D_i_W_emb = nn.Parameter(
+ torch.empty(7, model_args['time_emb_dim']))
+
+ # Decoupled Spatial Temporal Layer
+ self.layers = nn.ModuleList([DecoupleLayer(
+ self._hidden_dim, fk_dim=self._forecast_dim, first=True, **model_args)])
+ for _ in range(self._num_layers - 1):
+ self.layers.append(DecoupleLayer(
+ self._hidden_dim, fk_dim=self._forecast_dim, **model_args))
+
+ # dynamic and static hidden graph constructor
+ if model_args['dy_graph']:
+ self.dynamic_graph_constructor = DynamicGraphConstructor(
+ **model_args)
+
+ # node embeddings
+ self.node_emb_u = nn.Parameter(
+ torch.empty(self._num_nodes, self._node_dim))
+ self.node_emb_d = nn.Parameter(
+ torch.empty(self._num_nodes, self._node_dim))
+
+ # output layer
+ self.out_fc_1 = nn.Linear(self._forecast_dim, self._output_hidden)
+ self.out_fc_2 = nn.Linear(self._output_hidden, model_args['gap'])
+
+ self.reset_parameter()
+
+ def reset_parameter(self):
+ nn.init.xavier_uniform_(self.node_emb_u)
+ nn.init.xavier_uniform_(self.node_emb_d)
+ nn.init.xavier_uniform_(self.T_i_D_emb)
+ nn.init.xavier_uniform_(self.D_i_W_emb)
+
+ def _graph_constructor(self, **inputs):
+ E_d = inputs['E_d']
+ E_u = inputs['E_u']
+ if self._model_args['sta_graph']:
+ static_graph = [F.softmax(F.relu(torch.mm(E_d, E_u.T)), dim=1)]
+ else:
+ static_graph = []
+ if self._model_args['dy_graph']:
+ dynamic_graph = self.dynamic_graph_constructor(**inputs)
+ else:
+ dynamic_graph = []
+ return static_graph, dynamic_graph
+
+ def _prepare_inputs(self, X):
+ num_feat = self._model_args['num_feat']
+ # node embeddings
+ node_emb_u = self.node_emb_u # [N, d]
+ node_emb_d = self.node_emb_d # [N, d]
+ # time slot embedding
+ # [B, L, N, d]
+ T_i_D = self.T_i_D_emb[(X[:, :, :, num_feat] *
+ 288).type(torch.LongTensor)]
+ # [B, L, N, d]
+ D_i_W = self.D_i_W_emb[(X[:, :, :, num_feat+1]).type(torch.LongTensor)]
+ # traffic signals
+ X = X[:, :, :, :num_feat]
+
+ return X, node_emb_u, node_emb_d, T_i_D, D_i_W
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool, **kwargs) -> torch.Tensor:
+ """
+
+ Args:
+ history_data (Tensor): Input data with shape: [B, L, N, C]
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L, N, C]
+ """
+
+ X = history_data
+ # ==================== Prepare Input Data ==================== #
+ X, E_u, E_d, T_D, D_W = self._prepare_inputs(X)
+
+ # ========================= Construct Graphs ========================== #
+ static_graph, dynamic_graph = self._graph_constructor(
+ E_u=E_u, E_d=E_d, X=X, T_D=T_D, D_W=D_W)
+
+ # Start embedding layer
+ X = self.embedding(X)
+
+ spa_forecast_hidden_list = []
+ tem_forecast_hidden_list = []
+
+ tem_backcast_seq_res = X
+ for index, layer in enumerate(self.layers):
+ tem_backcast_seq_res, spa_forecast_hidden, tem_forecast_hidden = layer(
+ tem_backcast_seq_res, dynamic_graph, static_graph, E_u, E_d, T_D, D_W)
+ spa_forecast_hidden_list.append(spa_forecast_hidden)
+ tem_forecast_hidden_list.append(tem_forecast_hidden)
+
+ # Output Layer
+ spa_forecast_hidden = sum(spa_forecast_hidden_list)
+ tem_forecast_hidden = sum(tem_forecast_hidden_list)
+ forecast_hidden = spa_forecast_hidden + tem_forecast_hidden
+
+ # regression layer
+ forecast = self.out_fc_2(
+ F.relu(self.out_fc_1(F.relu(forecast_hidden))))
+ forecast = forecast.transpose(1, 2).contiguous().view(
+ forecast.shape[0], forecast.shape[2], -1)
+
+ # reshape
+ forecast = forecast.transpose(1, 2).unsqueeze(-1)
+ return forecast
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/decouple/estimation_gate.py b/basicts/archs/arch_zoo/d2stgnn_arch/decouple/estimation_gate.py
new file mode 100644
index 0000000..0789386
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/decouple/estimation_gate.py
@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+
+class EstimationGate(nn.Module):
+ r"""
+ The spatial gate module.
+ """
+ def __init__(self, node_emb_dim, time_emb_dim, hidden_dim, input_seq_len):
+ super().__init__()
+ self.FC1 = nn.Linear(2 * node_emb_dim + time_emb_dim * 2, hidden_dim)
+ self.act = nn.ReLU()
+ self.FC2 = nn.Linear(hidden_dim, 1)
+
+ def forward(self, node_embedding1, node_embedding2, T_D, D_W, X):
+ B, L, N, D = T_D.shape
+ spatial_gate_feat = torch.cat([T_D, D_W, node_embedding1.unsqueeze(0).unsqueeze(0).expand(B, L, -1, -1), node_embedding2.unsqueeze(0).unsqueeze(0).expand(B, L, -1, -1)], dim=-1)
+ hidden = self.FC1(spatial_gate_feat)
+ hidden = self.act(hidden)
+ # activation
+ spatial_gate = torch.sigmoid(self.FC2(hidden))[:, -X.shape[1]:, :, :]
+ X = X * spatial_gate
+ return X
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/decouple/residual_decomp.py b/basicts/archs/arch_zoo/d2stgnn_arch/decouple/residual_decomp.py
new file mode 100644
index 0000000..186c4ff
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/decouple/residual_decomp.py
@@ -0,0 +1,15 @@
+import torch.nn as nn
+
+class ResidualDecomp(nn.Module):
+ r"""
+ Residual decomposition.
+ """
+ def __init__(self, input_shape):
+ super().__init__()
+ self.ln = nn.LayerNorm(input_shape[-1])
+ self.ac = nn.ReLU()
+
+ def forward(self, x, y):
+ u = x - self.ac(y)
+ u = self.ln(u)
+ return u
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/__init__.py b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/__init__.py
new file mode 100644
index 0000000..9400012
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/__init__.py
@@ -0,0 +1 @@
+from ..difusion_block.dif_block import DifBlock
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_block.py b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_block.py
new file mode 100644
index 0000000..c360f3c
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_block.py
@@ -0,0 +1,34 @@
+import torch.nn as nn
+
+from ..decouple.residual_decomp import ResidualDecomp
+from .forecast import Forecast
+from .dif_model import STLocalizedConv
+
+
+class DifBlock(nn.Module):
+ def __init__(self, hidden_dim, fk_dim=256, use_pre=None, dy_graph=None, sta_graph=None, **model_args):
+ super().__init__()
+ self.pre_defined_graph = model_args['adjs']
+ self.localized_st_conv = STLocalizedConv(hidden_dim, pre_defined_graph=self.pre_defined_graph, \
+ use_pre=use_pre, dy_graph=dy_graph, sta_graph=sta_graph, **model_args)
+ # sub and norm
+ self.residual_decompose = ResidualDecomp([-1, -1, -1, hidden_dim])
+ # forecast
+ self.forecast_branch = Forecast(
+ hidden_dim, fk_dim=fk_dim, **model_args)
+ # backcast
+ self.backcast_branch = nn.Linear(hidden_dim, hidden_dim)
+
+ def forward(self, X, X_spa, dynamic_graph, static_graph):
+ Z = self.localized_st_conv(X_spa, dynamic_graph, static_graph)
+ # forecast branch
+ forecast_hidden = self.forecast_branch(
+ X_spa, Z, self.localized_st_conv, dynamic_graph, static_graph)
+ # backcast branch
+ backcast_seq = self.backcast_branch(Z)
+ # Residual Decomposition
+ backcast_seq = backcast_seq
+ X = X[:, -backcast_seq.shape[1]:, :, :]
+ backcast_seq_res = self.residual_decompose(X, backcast_seq)
+
+ return backcast_seq_res, forecast_hidden
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_model.py b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_model.py
new file mode 100644
index 0000000..dff8b70
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/dif_model.py
@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+
+
+class STLocalizedConv(nn.Module):
+ def __init__(self, hidden_dim, pre_defined_graph=None, use_pre=None, dy_graph=None, sta_graph=None, **model_args):
+ super().__init__()
+ # gated temporal conv
+ self.k_s = model_args['k_s']
+ self.k_t = model_args['k_t']
+ self.hidden_dim = hidden_dim
+
+ # graph conv
+ self.pre_defined_graph = pre_defined_graph
+ self.use_predefined_graph = use_pre
+ self.use_dynamic_hidden_graph = dy_graph
+ self.use_static__hidden_graph = sta_graph
+
+ self.support_len = len(self.pre_defined_graph) + \
+ int(dy_graph) + int(sta_graph)
+ self.num_matric = (int(use_pre) * len(self.pre_defined_graph) + len(
+ self.pre_defined_graph) * int(dy_graph) + int(sta_graph)) * self.k_s + 1
+ self.dropout = nn.Dropout(model_args['dropout'])
+ self.pre_defined_graph = self.get_graph(self.pre_defined_graph)
+
+ self.fc_list_updt = nn.Linear(
+ self.k_t * hidden_dim, self.k_t * hidden_dim, bias=False)
+ self.gcn_updt = nn.Linear(
+ self.hidden_dim*self.num_matric, self.hidden_dim)
+
+ # others
+ self.bn = nn.BatchNorm2d(self.hidden_dim)
+ self.activation = nn.ReLU()
+
+ def gconv(self, support, X_k, X_0):
+ out = [X_0]
+ for graph in support:
+ if len(graph.shape) == 2: # staitic or predefined graph
+ pass
+ else:
+ graph = graph.unsqueeze(1)
+ H_k = torch.matmul(graph, X_k)
+ out.append(H_k)
+ out = torch.cat(out, dim=-1)
+ out = self.gcn_updt(out)
+ out = self.dropout(out)
+ return out
+
+ def get_graph(self, support):
+ # Only used in static including static hidden graph and predefined graph, but not used for dynamic graph.
+ graph_ordered = []
+ mask = 1 - torch.eye(support[0].shape[0]).to(support[0].device)
+ for graph in support:
+ k_1_order = graph # 1 order
+ graph_ordered.append(k_1_order * mask)
+ # e.g., order = 3, k=[2, 3]; order = 2, k=[2]
+ for k in range(2, self.k_s+1):
+ k_1_order = torch.matmul(graph, k_1_order)
+ graph_ordered.append(k_1_order * mask)
+ # get st localed graph
+ st_local_graph = []
+ for graph in graph_ordered:
+ graph = graph.unsqueeze(-2).expand(-1, self.k_t, -1)
+ graph = graph.reshape(
+ graph.shape[0], graph.shape[1] * graph.shape[2])
+ # [num_nodes, kernel_size x num_nodes]
+ st_local_graph.append(graph)
+ # [order, num_nodes, kernel_size x num_nodes]
+ return st_local_graph
+
+ def forward(self, X, dynamic_graph, static_graph):
+ # X: [bs, seq, nodes, feat]
+ # [bs, seq, num_nodes, ks, num_feat]
+ X = X.unfold(1, self.k_t, 1).permute(0, 1, 2, 4, 3)
+ # seq_len is changing
+ batch_size, seq_len, num_nodes, kernel_size, num_feat = X.shape
+
+ # support
+ support = []
+ # predefined graph
+ if self.use_predefined_graph:
+ support = support + self.pre_defined_graph
+ # dynamic graph
+ if self.use_dynamic_hidden_graph:
+ # k_order is caled in dynamic_graph_constructor component
+ support = support + dynamic_graph
+ # predefined graphs and static hidden graphs
+ if self.use_static__hidden_graph:
+ support = support + self.get_graph(static_graph)
+
+ # parallelize
+ X = X.reshape(batch_size, seq_len, num_nodes, kernel_size * num_feat)
+ # batch_size, seq_len, num_nodes, kernel_size * hidden_dim
+ out = self.fc_list_updt(X)
+ out = self.activation(out)
+ out = out.view(batch_size, seq_len, num_nodes, kernel_size, num_feat)
+ X_0 = torch.mean(out, dim=-2)
+ # batch_size, seq_len, kernel_size x num_nodes, hidden_dim
+ X_k = out.transpose(-3, -2).reshape(batch_size,
+ seq_len, kernel_size*num_nodes, num_feat)
+ # Nx3N 3NxD -> NxD: batch_size, seq_len, num_nodes, hidden_dim
+ hidden = self.gconv(support, X_k, X_0)
+ return hidden
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/forecast.py b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/forecast.py
new file mode 100644
index 0000000..fa05190
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/difusion_block/forecast.py
@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+
+class Forecast(nn.Module):
+ def __init__(self, hidden_dim, fk_dim=None, **model_args):
+ super().__init__()
+ self.k_t = model_args['k_t']
+ self.output_seq_len = model_args['seq_length']
+ self.forecast_fc = nn.Linear(hidden_dim, fk_dim)
+ self.model_args = model_args
+
+ def forward(self, X, H, st_l_conv, dynamic_graph, static_graph):
+ [B, seq_len_remain, B, D] = H.shape
+ [B, seq_len_input, B, D] = X.shape
+
+ predict = []
+ history = X
+ predict.append(H[:, -1, :, :].unsqueeze(1))
+ for _ in range(int(self.output_seq_len / self.model_args['gap'])-1):
+ _1 = predict[-self.k_t:]
+ if len(_1) < self.k_t:
+ sub = self.k_t - len(_1)
+ _2 = history[:, -sub:, :, :]
+ _1 = torch.cat([_2] + _1, dim=1)
+ else:
+ _1 = torch.cat(_1, dim=1)
+ predict.append(st_l_conv(_1, dynamic_graph, static_graph))
+ predict = torch.cat(predict, dim=1)
+ predict = self.forecast_fc(predict)
+ return predict
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/dy_graph_conv.py b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/dy_graph_conv.py
new file mode 100644
index 0000000..fe7c585
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/dy_graph_conv.py
@@ -0,0 +1,50 @@
+import torch.nn as nn
+
+from .utils import *
+
+
+class DynamicGraphConstructor(nn.Module):
+ def __init__(self, **model_args):
+ super().__init__()
+ # model args
+ self.k_s = model_args['k_s'] # spatial order
+ self.k_t = model_args['k_t'] # temporal kernel size
+ # hidden dimension of
+ self.hidden_dim = model_args['num_hidden']
+ # trainable node embedding dimension
+ self.node_dim = model_args['node_hidden']
+
+ self.distance_function = DistanceFunction(**model_args)
+ self.mask = Mask(**model_args)
+ self.normalizer = Normalizer()
+ self.multi_order = MultiOrder(order=self.k_s)
+
+ def st_localization(self, graph_ordered):
+ st_local_graph = []
+ for modality_i in graph_ordered:
+ for k_order_graph in modality_i:
+ k_order_graph = k_order_graph.unsqueeze(
+ -2).expand(-1, -1, self.k_t, -1)
+ k_order_graph = k_order_graph.reshape(
+ k_order_graph.shape[0], k_order_graph.shape[1], k_order_graph.shape[2] * k_order_graph.shape[3])
+ st_local_graph.append(k_order_graph)
+ return st_local_graph
+
+ def forward(self, **inputs):
+ X = inputs['X']
+ E_d = inputs['E_d']
+ E_u = inputs['E_u']
+ T_D = inputs['T_D']
+ D_W = inputs['D_W']
+ # distance calculation
+ dist_mx = self.distance_function(X, E_d, E_u, T_D, D_W)
+ # mask
+ dist_mx = self.mask(dist_mx)
+ # normalization
+ dist_mx = self.normalizer(dist_mx)
+ # multi order
+ mul_mx = self.multi_order(dist_mx)
+ # spatial temporal localization
+ dynamic_graphs = self.st_localization(mul_mx)
+
+ return dynamic_graphs
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/__init__.py b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/__init__.py
new file mode 100644
index 0000000..f3c9349
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/__init__.py
@@ -0,0 +1,3 @@
+from .mask import *
+from .normalizer import *
+from .distance import *
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/distance.py b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/distance.py
new file mode 100644
index 0000000..ce1e642
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/distance.py
@@ -0,0 +1,57 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DistanceFunction(nn.Module):
+ def __init__(self, **model_args):
+ super().__init__()
+ # attributes
+ self.hidden_dim = model_args['num_hidden']
+ self.node_dim = model_args['node_hidden']
+ self.time_slot_emb_dim = self.hidden_dim
+ self.input_seq_len = model_args['seq_length']
+ # Time Series Feature Extraction
+ self.dropout = nn.Dropout(model_args['dropout'])
+ self.fc_ts_emb1 = nn.Linear(self.input_seq_len, self.hidden_dim * 2)
+ self.fc_ts_emb2 = nn.Linear(self.hidden_dim * 2, self.hidden_dim)
+ self.ts_feat_dim= self.hidden_dim
+ # Time Slot Embedding Extraction
+ self.time_slot_embedding = nn.Linear(model_args['time_emb_dim'], self.time_slot_emb_dim)
+ # Distance Score
+ self.all_feat_dim = self.ts_feat_dim + self.node_dim + model_args['time_emb_dim']*2
+ self.WQ = nn.Linear(self.all_feat_dim, self.hidden_dim, bias=False)
+ self.WK = nn.Linear(self.all_feat_dim, self.hidden_dim, bias=False)
+ self.bn = nn.BatchNorm1d(self.hidden_dim*2)
+
+ def reset_parameters(self):
+ for q_vec in self.q_vecs:
+ nn.init.xavier_normal_(q_vec.data)
+ for bias in self.biases:
+ nn.init.zeros_(bias.data)
+
+ def forward(self, X, E_d, E_u, T_D, D_W):
+ # last pooling
+ T_D = T_D[:, -1, :, :]
+ D_W = D_W[:, -1, :, :]
+ # dynamic information
+ X = X[:, :, :, 0].transpose(1, 2).contiguous() # X->[batch_size, seq_len, num_nodes]->[batch_size, num_nodes, seq_len]
+ [batch_size, num_nodes, seq_len] = X.shape
+ X = X.view(batch_size * num_nodes, seq_len)
+ dy_feat = self.fc_ts_emb2(self.dropout(self.bn(F.relu(self.fc_ts_emb1(X))))) # [batchsize, num_nodes, hidden_dim]
+ dy_feat = dy_feat.view(batch_size, num_nodes, -1)
+ # node embedding
+ emb1 = E_d.unsqueeze(0).expand(batch_size, -1, -1)
+ emb2 = E_u.unsqueeze(0).expand(batch_size, -1, -1)
+ # distance calculation
+ X1 = torch.cat([dy_feat, T_D, D_W, emb1], dim=-1) # hidden state for calculating distance
+ X2 = torch.cat([dy_feat, T_D, D_W, emb2], dim=-1) # hidden state for calculating distance
+ X = [X1, X2]
+ adjacent_list = []
+ for _ in X:
+ Q = self.WQ(_)
+ K = self.WK(_)
+ QKT = torch.bmm(Q, K.transpose(-1, -2)) / math.sqrt(self.hidden_dim)
+ W = torch.softmax(QKT, dim=-1)
+ adjacent_list.append(W)
+ return adjacent_list
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/mask.py b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/mask.py
new file mode 100644
index 0000000..1f3b6ec
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/mask.py
@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+
+class Mask(nn.Module):
+ def __init__(self, **model_args):
+ super().__init__()
+ self.mask = model_args['adjs']
+
+ def _mask(self, index, adj):
+ mask = self.mask[index] + torch.ones_like(self.mask[index]) * 1e-7
+ return mask.to(adj.device) * adj
+
+ def forward(self, adj):
+ result = []
+ for index, _ in enumerate(adj):
+ result.append(self._mask(index, _))
+ return result
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/normalizer.py b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/normalizer.py
new file mode 100644
index 0000000..efab418
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/dynamic_graph_conv/utils/normalizer.py
@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+from basicts.utils.misc import remove_nan_inf
+
+class Normalizer(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def _norm(self, graph):
+ degree = torch.sum(graph, dim=2)
+ degree = remove_nan_inf(1 / degree)
+ degree = torch.diag_embed(degree)
+ P = torch.bmm(degree, graph)
+ return P
+
+ def forward(self, adj):
+ return [self._norm(_) for _ in adj]
+
+class MultiOrder(nn.Module):
+ def __init__(self, order=2):
+ super().__init__()
+ self.order = order
+
+ def _multi_order(self, graph):
+ graph_ordered = []
+ k_1_order = graph # 1 order
+ mask = torch.eye(graph.shape[1]).to(graph.device)
+ mask = 1 - mask
+ graph_ordered.append(k_1_order * mask)
+ for k in range(2, self.order+1): # e.g., order = 3, k=[2, 3]; order = 2, k=[2]
+ k_1_order = torch.matmul(k_1_order, graph)
+ graph_ordered.append(k_1_order * mask)
+ return graph_ordered
+
+ def forward(self, adj):
+ return [self._multi_order(_) for _ in adj]
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/__init__.py b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/__init__.py
new file mode 100644
index 0000000..521c8c9
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/__init__.py
@@ -0,0 +1 @@
+from .inh_block import InhBlock
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/forecast.py b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/forecast.py
new file mode 100644
index 0000000..e9224a4
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/forecast.py
@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+
+
+class Forecast(nn.Module):
+ def __init__(self, hidden_dim, fk_dim, **model_args):
+ super().__init__()
+ self.output_seq_len = model_args['seq_length']
+ self.model_args = model_args
+
+ self.forecast_fc = nn.Linear(hidden_dim, fk_dim)
+
+ def forward(self, X, RNN_H, Z, transformer_layer, rnn_layer, pe):
+ [B, L, N, D] = X.shape
+ [L, B_N, D] = RNN_H.shape
+ [L, B_N, D] = Z.shape
+
+ predict = [Z[-1, :, :].unsqueeze(0)]
+ for _ in range(int(self.output_seq_len / self.model_args['gap'])-1):
+ # RNN
+ _gru = rnn_layer.gru_cell(predict[-1][0], RNN_H[-1]).unsqueeze(0)
+ RNN_H = torch.cat([RNN_H, _gru], dim=0)
+ # Positional Encoding
+ if pe is not None:
+ RNN_H = pe(RNN_H)
+ # Transformer
+ _Z = transformer_layer(_gru, K=RNN_H, V=RNN_H)
+ predict.append(_Z)
+
+ predict = torch.cat(predict, dim=0)
+ predict = predict.reshape(-1, B, N, D)
+ predict = predict.transpose(0, 1)
+ predict = self.forecast_fc(predict)
+ return predict
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_block.py b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_block.py
new file mode 100644
index 0000000..431e1e2
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_block.py
@@ -0,0 +1,72 @@
+import math
+import torch
+import torch.nn as nn
+
+from ..decouple.residual_decomp import ResidualDecomp
+from .inh_model import RNNLayer, TransformerLayer
+from .forecast import Forecast
+
+
+class PositionalEncoding(nn.Module):
+ def __init__(self, d_model, dropout=None, max_len: int = 5000):
+ super().__init__()
+ self.dropout = nn.Dropout(p=dropout)
+ position = torch.arange(max_len).unsqueeze(1)
+ div_term = torch.exp(torch.arange(0, d_model, 2)
+ * (-math.log(10000.0) / d_model))
+ pe = torch.zeros(max_len, 1, d_model)
+ pe[:, 0, 0::2] = torch.sin(position * div_term)
+ pe[:, 0, 1::2] = torch.cos(position * div_term)
+ self.register_buffer('pe', pe)
+
+ def forward(self, X):
+ X = X + self.pe[:X.size(0)]
+ X = self.dropout(X)
+ return X
+
+
+class InhBlock(nn.Module):
+ def __init__(self, hidden_dim, num_heads=4, bias=True, fk_dim=256, first=None, **model_args):
+ super().__init__()
+ self.num_feat = hidden_dim
+ self.hidden_dim = hidden_dim
+
+ if first:
+ self.pos_encoder = PositionalEncoding(
+ hidden_dim, model_args['dropout'])
+ else:
+ self.pos_encoder = None
+ self.rnn_layer = RNNLayer(hidden_dim, model_args['dropout'])
+ self.transformer_layer = TransformerLayer(
+ hidden_dim, num_heads, model_args['dropout'], bias)
+ # forecast
+ self.forecast_block = Forecast(hidden_dim, fk_dim, **model_args)
+ # backcast
+ self.backcast_fc = nn.Linear(hidden_dim, hidden_dim)
+ # sub residual
+ self.sub_and_norm = ResidualDecomp([-1, -1, -1, hidden_dim])
+
+ def forward(self, X):
+ [batch_size, seq_len, num_nodes, num_feat] = X.shape
+ # Temporal Model
+ # RNN
+ RNN_H_raw = self.rnn_layer(X)
+ # Positional Encoding
+ if self.pos_encoder is not None:
+ RNN_H = self.pos_encoder(RNN_H_raw)
+ else:
+ RNN_H = RNN_H_raw
+ # MultiHead Self Attention
+ Z = self.transformer_layer(RNN_H, RNN_H, RNN_H)
+
+ # forecast branch
+ forecast_hidden = self.forecast_block(
+ X, RNN_H_raw, Z, self.transformer_layer, self.rnn_layer, self.pos_encoder)
+
+ # backcast branch
+ Z = Z.reshape(seq_len, batch_size, num_nodes, num_feat)
+ Z = Z.transpose(0, 1)
+ backcast_seq = self.backcast_fc(Z)
+ backcast_seq_res = self.sub_and_norm(X, backcast_seq)
+
+ return backcast_seq_res, forecast_hidden
diff --git a/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_model.py b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_model.py
new file mode 100644
index 0000000..ce8e279
--- /dev/null
+++ b/basicts/archs/arch_zoo/d2stgnn_arch/inherent_block/inh_model.py
@@ -0,0 +1,37 @@
+import torch as th
+import torch.nn as nn
+from torch.nn import MultiheadAttention
+
+
+class RNNLayer(nn.Module):
+ def __init__(self, hidden_dim, dropout=None):
+ super().__init__()
+ self.hidden_dim = hidden_dim
+ self.gru_cell = nn.GRUCell(hidden_dim, hidden_dim)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, X):
+ [batch_size, seq_len, num_nodes, hidden_dim] = X.shape
+ X = X.transpose(1, 2).reshape(
+ batch_size * num_nodes, seq_len, hidden_dim)
+ hx = th.zeros_like(X[:, 0, :])
+ output = []
+ for _ in range(X.shape[1]):
+ hx = self.gru_cell(X[:, _, :], hx)
+ output.append(hx)
+ output = th.stack(output, dim=0)
+ output = self.dropout(output)
+ return output
+
+
+class TransformerLayer(nn.Module):
+ def __init__(self, hidden_dim, num_heads=4, dropout=None, bias=True):
+ super().__init__()
+ self.multi_head_self_attention = MultiheadAttention(
+ hidden_dim, num_heads, dropout=dropout, bias=bias)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, X, K, V):
+ Z = self.multi_head_self_attention(X, K, V)[0]
+ Z = self.dropout(Z)
+ return Z
diff --git a/basicts/archs/arch_zoo/dcrnn_arch/__init__.py b/basicts/archs/arch_zoo/dcrnn_arch/__init__.py
new file mode 100644
index 0000000..0307d9f
--- /dev/null
+++ b/basicts/archs/arch_zoo/dcrnn_arch/__init__.py
@@ -0,0 +1,3 @@
+from .dcrnn_arch import DCRNN
+
+__all__ = ['DCRNN']
diff --git a/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_arch.py b/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_arch.py
new file mode 100644
index 0000000..759be73
--- /dev/null
+++ b/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_arch.py
@@ -0,0 +1,168 @@
+import torch
+from torch import nn
+import numpy as np
+
+from .dcrnn_cell import DCGRUCell
+
+
+def count_parameters(model):
+ return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class Seq2SeqAttrs:
+ def __init__(self, adj_mx, **model_kwargs):
+ self.adj_mx = adj_mx
+ self.max_diffusion_step = int(
+ model_kwargs.get("max_diffusion_step", 2))
+ self.cl_decay_steps = int(model_kwargs.get("cl_decay_steps", 1000))
+ self.filter_type = model_kwargs.get("filter_type", "laplacian")
+ self.num_nodes = int(model_kwargs.get("num_nodes", 1))
+ self.num_rnn_layers = int(model_kwargs.get("num_rnn_layers", 1))
+ self.rnn_units = int(model_kwargs.get("rnn_units"))
+ self.hidden_state_size = self.num_nodes * self.rnn_units
+
+
+class EncoderModel(nn.Module, Seq2SeqAttrs):
+ def __init__(self, adj_mx, **model_kwargs):
+ nn.Module.__init__(self)
+ Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
+ self.input_dim = int(model_kwargs.get("input_dim", 1))
+ self.seq_len = int(model_kwargs.get("seq_len")) # for the encoder
+ self.dcgru_layers = nn.ModuleList(
+ [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes) for _ in range(self.num_rnn_layers)])
+
+ def forward(self, inputs, hidden_state=None):
+ batch_size, _ = inputs.size()
+ if hidden_state is None:
+ hidden_state = torch.zeros(
+ (self.num_rnn_layers, batch_size, self.hidden_state_size)).to(inputs.device)
+ hidden_states = []
+ output = inputs
+ for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
+ next_hidden_state = dcgru_layer(output, hidden_state[layer_num])
+ hidden_states.append(next_hidden_state)
+ output = next_hidden_state
+
+ # runs in O(num_layers) so not too slow
+ return output, torch.stack(hidden_states)
+
+
+class DecoderModel(nn.Module, Seq2SeqAttrs):
+ def __init__(self, adj_mx, **model_kwargs):
+ # super().__init__(is_training, adj_mx, **model_kwargs)
+ nn.Module.__init__(self)
+ Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
+ self.output_dim = int(model_kwargs.get("output_dim", 1))
+ self.horizon = int(model_kwargs.get("horizon", 1)) # for the decoder
+ self.projection_layer = nn.Linear(self.rnn_units, self.output_dim)
+ self.dcgru_layers = nn.ModuleList(
+ [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes) for _ in range(self.num_rnn_layers)])
+
+ def forward(self, inputs, hidden_state=None):
+ hidden_states = []
+ output = inputs
+ for layer_num, dcgru_layer in enumerate(self.dcgru_layers):
+ next_hidden_state = dcgru_layer(output, hidden_state[layer_num])
+ hidden_states.append(next_hidden_state)
+ output = next_hidden_state
+
+ projected = self.projection_layer(output.view(-1, self.rnn_units))
+ output = projected.view(-1, self.num_nodes * self.output_dim)
+
+ return output, torch.stack(hidden_states)
+
+
+class DCRNN(nn.Module, Seq2SeqAttrs):
+ """
+ Paper: Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting
+ Link: https://arxiv.org/abs/1707.01926
+ Codes are modified from the official repo:
+ https://github.com/chnsh/DCRNN_PyTorch/blob/pytorch_scratch/model/pytorch/dcrnn_cell.py,
+ https://github.com/chnsh/DCRNN_PyTorch/blob/pytorch_scratch/model/pytorch/dcrnn_model.py
+ """
+
+ def __init__(self, adj_mx, **model_kwargs):
+ super().__init__()
+ Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs)
+ self.encoder_model = EncoderModel(adj_mx, **model_kwargs)
+ self.decoder_model = DecoderModel(adj_mx, **model_kwargs)
+ self.cl_decay_steps = int(model_kwargs.get("cl_decay_steps", 2000))
+ self.use_curriculum_learning = bool(
+ model_kwargs.get("use_curriculum_learning", False))
+
+ def _compute_sampling_threshold(self, batches_seen):
+ return self.cl_decay_steps / (
+ self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps))
+
+ def encoder(self, inputs):
+ encoder_hidden_state = None
+ for t in range(self.encoder_model.seq_len):
+ _, encoder_hidden_state = self.encoder_model(
+ inputs[t], encoder_hidden_state)
+
+ return encoder_hidden_state
+
+ def decoder(self, encoder_hidden_state, labels=None, batches_seen=None):
+ batch_size = encoder_hidden_state.size(1)
+ go_symbol = torch.zeros(
+ (batch_size, self.num_nodes * self.decoder_model.output_dim)).to(encoder_hidden_state.device)
+ decoder_hidden_state = encoder_hidden_state
+ decoder_input = go_symbol
+
+ outputs = []
+
+ for t in range(self.decoder_model.horizon):
+ decoder_output, decoder_hidden_state = self.decoder_model(
+ decoder_input, decoder_hidden_state)
+ decoder_input = decoder_output
+ outputs.append(decoder_output)
+ if self.training and self.use_curriculum_learning:
+ c = np.random.uniform(0, 1)
+ if c < self._compute_sampling_threshold(batches_seen):
+ decoder_input = labels[t]
+ outputs = torch.stack(outputs)
+ return outputs
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor = None, batch_seen: int = None, **kwargs) -> torch.Tensor:
+ """Feedforward function of DCRNN.
+
+ Args:
+ history_data (torch.Tensor): history data with shape [L, B, N*C]
+ future_data (torch.Tensor, optional): future data with shape [L, B, N*C_out]
+ batch_seen (int, optional): batches seen till now, used for curriculum learning. Defaults to None.
+
+ Returns:
+ torch.Tensor: prediction wit shape [L, B, N*C_out]
+ """
+
+ # reshape data
+ batch_size, length, num_nodes, channels = history_data.shape
+ history_data = history_data.reshape(batch_size, length, num_nodes * channels) # [B, L, N*C]
+ history_data = history_data.transpose(0, 1) # [L, B, N*C]
+
+ if future_data is not None:
+ batch_size, length, num_nodes, channels = future_data.shape
+ future_data = future_data.reshape(batch_size, length, num_nodes * channels) # [B, L, N*C]
+ future_data = future_data.transpose(0, 1) # [L, B, N*C]
+
+ # DCRNN
+ encoder_hidden_state = self.encoder(history_data)
+ outputs = self.decoder(encoder_hidden_state, future_data,
+ batches_seen=batch_seen) # [L, B, N*C_out]
+
+ # reshape to B, L, N, C
+ L, B, _ = outputs.shape
+ outputs = outputs.transpose(0, 1) # [B, L, N*C_out]
+ outputs = outputs.view(B, L, self.num_nodes,
+ self.decoder_model.output_dim)
+
+ if not self.training:
+ assert future_data == None, "Future data should not be visible when validating/testing."
+ else:
+ pass
+
+ if batch_seen == 0:
+ print("Warning: decoder only takes the first dimension as groundtruth.")
+ print("Parameter Number: ".format(count_parameters(self)))
+ print(count_parameters(self))
+ return outputs
diff --git a/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_cell.py b/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_cell.py
new file mode 100644
index 0000000..4453fb1
--- /dev/null
+++ b/basicts/archs/arch_zoo/dcrnn_arch/dcrnn_cell.py
@@ -0,0 +1,140 @@
+import torch
+
+
+class LayerParams:
+ """Layer parameters."""
+
+ def __init__(self, rnn_network: torch.nn.Module, layer_type: str):
+ self._rnn_network = rnn_network
+ self._params_dict = {}
+ self._biases_dict = {}
+ self._type = layer_type
+
+ def get_weights(self, shape):
+ if shape not in self._params_dict:
+ nn_param = torch.nn.Parameter(torch.empty(*shape))
+ torch.nn.init.xavier_normal_(nn_param)
+ self._params_dict[shape] = nn_param
+ self._rnn_network.register_parameter(
+ '{}_weight_{}'.format(self._type, str(shape)),
+ nn_param)
+ return self._params_dict[shape]
+
+ def get_biases(self, length, bias_start=0.0):
+ if length not in self._biases_dict:
+ biases = torch.nn.Parameter(torch.empty(length))
+ torch.nn.init.constant_(biases, bias_start)
+ self._biases_dict[length] = biases
+ self._rnn_network.register_parameter(
+ '{}_biases_{}'.format(self._type, str(length)), biases)
+
+ return self._biases_dict[length]
+
+
+class DCGRUCell(torch.nn.Module):
+ """
+ Paper: Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting
+ Link: https://arxiv.org/abs/1707.01926
+ Codes are modified from the official repo:
+ https://github.com/chnsh/DCRNN_PyTorch/blob/pytorch_scratch/model/pytorch/dcrnn_cell.py,
+ https://github.com/chnsh/DCRNN_PyTorch/blob/pytorch_scratch/model/pytorch/dcrnn_model.py
+ Watch out the input groundtruth of decoder, which may cause bugs when you try to extend this code.
+ In order to train the model on multi-GPU, we send the parameter to different gpus in the feedforward process, which might hurt the efficiency.
+ """
+
+ def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, nonlinearity='tanh', use_gc_for_ru=True):
+ super().__init__()
+ self._activation = torch.tanh if nonlinearity == 'tanh' else torch.relu
+ # support other nonlinearities up here?
+ self._num_nodes = num_nodes
+ self._num_units = num_units
+ self._max_diffusion_step = max_diffusion_step
+ self._use_gc_for_ru = use_gc_for_ru
+ # for support in supports:
+ # self._supports.append(self._build_sparse_matrix(support))
+ self._supports = adj_mx
+
+ self._fc_params = LayerParams(self, 'fc')
+ self._gconv_params = LayerParams(self, 'gconv')
+
+ def forward(self, inputs, hx):
+ output_size = 2 * self._num_units
+ if self._use_gc_for_ru:
+ fn = self._gconv
+ else:
+ fn = self._fc
+ value = torch.sigmoid(fn(inputs, hx, output_size, bias_start=1.0))
+ value = torch.reshape(value, (-1, self._num_nodes, output_size))
+ r, u = torch.split(
+ tensor=value, split_size_or_sections=self._num_units, dim=-1)
+ r = torch.reshape(r, (-1, self._num_nodes * self._num_units))
+ u = torch.reshape(u, (-1, self._num_nodes * self._num_units))
+
+ c = self._gconv(inputs, r * hx, self._num_units)
+ if self._activation is not None:
+ c = self._activation(c)
+
+ new_state = u * hx + (1.0 - u) * c
+ return new_state
+
+ @staticmethod
+ def _concat(x, x_):
+ x_ = x_.unsqueeze(0)
+ return torch.cat([x, x_], dim=0)
+
+ def _fc(self, inputs, state, output_size, bias_start=0.0):
+ batch_size = inputs.shape[0]
+ inputs = torch.reshape(inputs, (batch_size * self._num_nodes, -1))
+ state = torch.reshape(state, (batch_size * self._num_nodes, -1))
+ inputs_and_state = torch.cat([inputs, state], dim=-1)
+ input_size = inputs_and_state.shape[-1]
+ weights = self._fc_params.get_weights(
+ (input_size, output_size)).to(inputs_and_state.device)
+ value = torch.sigmoid(torch.matmul(inputs_and_state, weights))
+ biases = self._fc_params.get_biases(output_size, bias_start)
+ value += biases.to(inputs_and_state.device)
+ return value
+
+ def _gconv(self, inputs, state, output_size, bias_start=0.0):
+ # Reshape input and state to (batch_size, num_nodes, input_dim/state_dim)
+ batch_size = inputs.shape[0]
+ inputs = torch.reshape(inputs, (batch_size, self._num_nodes, -1))
+ state = torch.reshape(state, (batch_size, self._num_nodes, -1))
+ inputs_and_state = torch.cat([inputs, state], dim=2)
+ input_size = inputs_and_state.size(2)
+
+ x = inputs_and_state
+ x0 = x.permute(1, 2, 0) # (num_nodes, total_arg_size, batch_size)
+ x0 = torch.reshape(
+ x0, shape=[self._num_nodes, input_size * batch_size])
+ x = torch.unsqueeze(x0, 0)
+
+ if self._max_diffusion_step == 0:
+ pass
+ else:
+ for support in self._supports:
+ x1 = torch.mm(support, x0)
+ x = self._concat(x, x1)
+
+ for k in range(2, self._max_diffusion_step + 1):
+ x2 = 2 * torch.mm(support, x1) - x0
+ x = self._concat(x, x2)
+ x1, x0 = x2, x1
+
+ # Adds for x itself.
+ num_matrices = len(self._supports) * self._max_diffusion_step + 1
+ x = torch.reshape(
+ x, shape=[num_matrices, self._num_nodes, input_size, batch_size])
+ x = x.permute(3, 1, 2, 0) # (batch_size, num_nodes, input_size, order)
+ x = torch.reshape(
+ x, shape=[batch_size * self._num_nodes, input_size * num_matrices])
+ weights = self._gconv_params.get_weights(
+ (input_size * num_matrices, output_size)).to(x.device)
+ # (batch_size * self._num_nodes, output_size)
+ x = torch.matmul(x, weights)
+
+ biases = self._gconv_params.get_biases(
+ output_size, bias_start).to(x.device)
+ x += biases
+ # Reshape res back to 2D: (batch_size, num_node, state_dim) -> (batch_size, num_node * state_dim)
+ return torch.reshape(x, [batch_size, self._num_nodes * output_size])
diff --git a/basicts/archs/arch_zoo/dgcrn_arch/__init__.py b/basicts/archs/arch_zoo/dgcrn_arch/__init__.py
new file mode 100644
index 0000000..8303eb4
--- /dev/null
+++ b/basicts/archs/arch_zoo/dgcrn_arch/__init__.py
@@ -0,0 +1,3 @@
+from .dgcrn_arch import DGCRN
+
+__all__ = ["DGCRN"]
diff --git a/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_arch.py b/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_arch.py
new file mode 100644
index 0000000..34ae93f
--- /dev/null
+++ b/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_arch.py
@@ -0,0 +1,247 @@
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+from .dgcrn_layer import *
+
+
+class DGCRN(nn.Module):
+ def __init__(self, gcn_depth, num_nodes, predefined_A=None, dropout=0.3, subgraph_size=20, node_dim=40, middle_dim=2, seq_length=12, in_dim=2, list_weight=[0.05, 0.95, 0.95], tanhalpha=3, cl_decay_steps=4000, rnn_size=64, hyperGNN_dim=16):
+ super(DGCRN, self).__init__()
+ self.output_dim = 1
+
+ self.num_nodes = num_nodes
+ self.dropout = dropout
+ self.predefined_A = predefined_A
+
+ self.seq_length = seq_length
+
+ self.emb1 = nn.Embedding(self.num_nodes, node_dim)
+ self.emb2 = nn.Embedding(self.num_nodes, node_dim)
+ self.lin1 = nn.Linear(node_dim, node_dim)
+ self.lin2 = nn.Linear(node_dim, node_dim)
+
+ self.idx = torch.arange(self.num_nodes)
+
+ self.rnn_size = rnn_size
+ self.in_dim = in_dim
+
+ self.hidden_size = self.rnn_size
+
+ dims_hyper = [self.hidden_size + in_dim,
+ hyperGNN_dim, middle_dim, node_dim]
+
+ self.GCN1_tg = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN2_tg = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN1_tg_de = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN2_tg_de = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN1_tg_1 = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN2_tg_1 = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN1_tg_de_1 = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.GCN2_tg_de_1 = gcn(dims_hyper, gcn_depth,
+ dropout, *list_weight, 'hyper')
+
+ self.fc_final = nn.Linear(self.hidden_size, self.output_dim)
+
+ self.alpha = tanhalpha
+ self.k = subgraph_size
+ dims = [in_dim + self.hidden_size, self.hidden_size]
+
+ self.gz1 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gz2 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gr1 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gr2 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gc1 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gc2 = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+
+ self.gz1_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gz2_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gr1_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gr2_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gc1_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+ self.gc2_de = gcn(dims, gcn_depth, dropout, *list_weight, 'RNN')
+
+ self.use_curriculum_learning = True
+ self.cl_decay_steps = cl_decay_steps
+ self.gcn_depth = gcn_depth
+
+ def preprocessing(self, adj, predefined_A):
+ adj = adj + torch.eye(self.num_nodes).to(adj.device)
+ adj = adj / torch.unsqueeze(adj.sum(-1), -1)
+ return [adj, predefined_A]
+
+ def step(self, input, Hidden_State, Cell_State, predefined_A, type='encoder', i=None):
+ x = input
+ x = x.transpose(1, 2).contiguous()
+
+ nodevec1 = self.emb1(self.idx)
+ nodevec2 = self.emb2(self.idx)
+
+ hyper_input = torch.cat(
+ (x, Hidden_State.view(-1, self.num_nodes, self.hidden_size)), 2)
+
+ if type == 'encoder':
+ filter1 = self.GCN1_tg(
+ hyper_input, predefined_A[0]) + self.GCN1_tg_1(hyper_input, predefined_A[1])
+ filter2 = self.GCN2_tg(
+ hyper_input, predefined_A[0]) + self.GCN2_tg_1(hyper_input, predefined_A[1])
+
+ if type == 'decoder':
+ filter1 = self.GCN1_tg_de(
+ hyper_input, predefined_A[0]) + self.GCN1_tg_de_1(hyper_input, predefined_A[1])
+ filter2 = self.GCN2_tg_de(
+ hyper_input, predefined_A[0]) + self.GCN2_tg_de_1(hyper_input, predefined_A[1])
+
+ nodevec1 = torch.tanh(self.alpha * torch.mul(nodevec1, filter1))
+ nodevec2 = torch.tanh(self.alpha * torch.mul(nodevec2, filter2))
+
+ a = torch.matmul(nodevec1, nodevec2.transpose(2, 1)) - \
+ torch.matmul(nodevec2, nodevec1.transpose(2, 1))
+
+ adj = F.relu(torch.tanh(self.alpha * a))
+
+ adp = self.preprocessing(adj, predefined_A[0])
+ adpT = self.preprocessing(adj.transpose(1, 2), predefined_A[1])
+
+ Hidden_State = Hidden_State.view(-1, self.num_nodes, self.hidden_size)
+ Cell_State = Cell_State.view(-1, self.num_nodes, self.hidden_size)
+
+ combined = torch.cat((x, Hidden_State), -1)
+
+ if type == 'encoder':
+ z = torch.sigmoid(self.gz1(combined, adp) +
+ self.gz2(combined, adpT))
+ r = torch.sigmoid(self.gr1(combined, adp) +
+ self.gr2(combined, adpT))
+
+ temp = torch.cat((x, torch.mul(r, Hidden_State)), -1)
+ Cell_State = torch.tanh(self.gc1(temp, adp) + self.gc2(temp, adpT))
+ elif type == 'decoder':
+ z = torch.sigmoid(self.gz1_de(combined, adp) +
+ self.gz2_de(combined, adpT))
+ r = torch.sigmoid(self.gr1_de(combined, adp) +
+ self.gr2_de(combined, adpT))
+
+ temp = torch.cat((x, torch.mul(r, Hidden_State)), -1)
+ Cell_State = torch.tanh(self.gc1_de(
+ temp, adp) + self.gc2_de(temp, adpT))
+
+ Hidden_State = torch.mul(z, Hidden_State) + \
+ torch.mul(1 - z, Cell_State)
+
+ return Hidden_State.view(-1, self.hidden_size), Cell_State.view(-1, self.hidden_size)
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool,**kwargs) -> torch.Tensor:
+ """Feedforward function of DGCRN.
+
+ Args:
+ history_data (torch.Tensor): historical data with shape [B, L, N, C].
+ future_data (torch.Tensor, optional): ground truth. Defaults to None.
+ batch_seen (int, optional): batch num. Defaults to None.
+ task_level (int, optional): curriculum learning level. Defaults to 12.
+
+ Returns:
+ torch.Tensor: prediction with shape [B, L, N, 1]
+ """
+ task_level = kwargs["task_level"]
+ input = history_data.transpose(1, 3)
+ ycl = future_data.transpose(1, 3)
+
+ self.idx = self.idx.to(input.device)
+
+ predefined_A = self.predefined_A
+ x = input
+
+ batch_size = x.size(0)
+ Hidden_State, Cell_State = self.initHidden(
+ batch_size * self.num_nodes, self.hidden_size)
+ Hidden_State = Hidden_State.to(input.device)
+ Cell_State = Cell_State.to(input.device)
+
+ outputs = None
+ for i in range(self.seq_length):
+ Hidden_State, Cell_State = self.step(
+ x[..., i].squeeze(-1), Hidden_State, Cell_State, predefined_A, 'encoder', i)
+
+ if outputs is None:
+ outputs = Hidden_State.unsqueeze(1)
+ else:
+ outputs = torch.cat((outputs, Hidden_State.unsqueeze(1)), 1)
+
+ go_symbol = torch.zeros(
+ (batch_size, self.output_dim, self.num_nodes)).to(input.device)
+ timeofday = ycl[:, [1], :, :]
+
+ decoder_input = go_symbol
+
+ outputs_final = []
+
+ for i in range(task_level):
+ try:
+ decoder_input = torch.cat(
+ [decoder_input, timeofday[..., i]], dim=1)
+ except:
+ print(decoder_input.shape, timeofday.shape)
+ sys.exit(0)
+ Hidden_State, Cell_State = self.step(
+ decoder_input, Hidden_State, Cell_State, predefined_A, 'decoder', None)
+
+ decoder_output = self.fc_final(Hidden_State)
+
+ decoder_input = decoder_output.view(
+ batch_size, self.num_nodes, self.output_dim).transpose(1, 2)
+ outputs_final.append(decoder_output)
+ if self.training and self.use_curriculum_learning:
+ c = np.random.uniform(0, 1)
+ if c < self._compute_sampling_threshold(batch_seen):
+ decoder_input = ycl[:, :1, :, i]
+
+ outputs_final = torch.stack(outputs_final, dim=1)
+
+ outputs_final = outputs_final.view(
+ batch_size, self.num_nodes, task_level, self.output_dim).transpose(1, 2)
+
+ ramdom_predict = torch.zeros(batch_size, self.seq_length - task_level,
+ self.num_nodes, self.output_dim).to(outputs_final.device)
+ outputs_final = torch.cat([outputs_final, ramdom_predict], dim=1)
+
+ return outputs_final
+
+ def initHidden(self, batch_size, hidden_size):
+ use_gpu = torch.cuda.is_available()
+ if use_gpu:
+ Hidden_State = Variable(
+ torch.zeros(batch_size, hidden_size))
+ Cell_State = Variable(
+ torch.zeros(batch_size, hidden_size))
+
+ nn.init.orthogonal_(Hidden_State)
+ nn.init.orthogonal_(Cell_State)
+
+ return Hidden_State, Cell_State
+ else:
+ Hidden_State = Variable(torch.zeros(batch_size, hidden_size))
+ Cell_State = Variable(torch.zeros(batch_size, hidden_size))
+ return Hidden_State, Cell_State
+
+ def _compute_sampling_threshold(self, batches_seen):
+ return self.cl_decay_steps / (
+ self.cl_decay_steps + np.exp(batches_seen / self.cl_decay_steps))
diff --git a/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_layer.py b/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_layer.py
new file mode 100644
index 0000000..29ecb31
--- /dev/null
+++ b/basicts/archs/arch_zoo/dgcrn_arch/dgcrn_layer.py
@@ -0,0 +1,68 @@
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class gconv_RNN(nn.Module):
+ def __init__(self):
+ super(gconv_RNN, self).__init__()
+
+ def forward(self, x, A):
+
+ x = torch.einsum('nvc,nvw->nwc', (x, A))
+ return x.contiguous()
+
+
+class gconv_hyper(nn.Module):
+ def __init__(self):
+ super(gconv_hyper, self).__init__()
+
+ def forward(self, x, A):
+ A = A.to(x.device)
+ x = torch.einsum('nvc,vw->nwc', (x, A))
+ return x.contiguous()
+
+
+class gcn(nn.Module):
+ def __init__(self, dims, gdep, dropout, alpha, beta, gamma, type=None):
+ super(gcn, self).__init__()
+ if type == 'RNN':
+ self.gconv = gconv_RNN()
+ self.gconv_preA = gconv_hyper()
+ self.mlp = nn.Linear((gdep + 1) * dims[0], dims[1])
+
+ elif type == 'hyper':
+ self.gconv = gconv_hyper()
+ self.mlp = nn.Sequential(
+ OrderedDict([('fc1', nn.Linear((gdep + 1) * dims[0], dims[1])),
+ ('sigmoid1', nn.Sigmoid()),
+ ('fc2', nn.Linear(dims[1], dims[2])),
+ ('sigmoid2', nn.Sigmoid()),
+ ('fc3', nn.Linear(dims[2], dims[3]))]))
+
+ self.gdep = gdep
+ self.alpha = alpha
+ self.beta = beta
+ self.gamma = gamma
+ self.type_GNN = type
+
+ def forward(self, x, adj):
+
+ h = x
+ out = [h]
+ if self.type_GNN == 'RNN':
+ for _ in range(self.gdep):
+ h = self.alpha * x + self.beta * self.gconv(
+ h, adj[0]) + self.gamma * self.gconv_preA(h, adj[1])
+ out.append(h)
+ else:
+ for _ in range(self.gdep):
+ h = self.alpha * x + self.gamma * self.gconv(h, adj)
+ out.append(h)
+ ho = torch.cat(out, dim=-1)
+
+ ho = self.mlp(ho)
+
+ return ho
diff --git a/basicts/archs/arch_zoo/fedformer_arch/__init__.py b/basicts/archs/arch_zoo/fedformer_arch/__init__.py
new file mode 100644
index 0000000..66615d0
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/__init__.py
@@ -0,0 +1,3 @@
+from .fedformer_arch import FEDformer
+
+__all__ = ["FEDformer"]
diff --git a/basicts/archs/arch_zoo/fedformer_arch/auto_correlation.py b/basicts/archs/arch_zoo/fedformer_arch/auto_correlation.py
new file mode 100644
index 0000000..a051191
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/auto_correlation.py
@@ -0,0 +1,245 @@
+import time
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.functional import interpolate
+
+
+def decor_time(func):
+ def func2(*args, **kw):
+ now = time.time()
+ y = func(*args, **kw)
+ t = time.time() - now
+ print('call <{}>, time={}'.format(func.__name__, t))
+ return y
+ return func2
+
+
+class AutoCorrelation(nn.Module):
+ """
+ AutoCorrelation Mechanism with the following two phases:
+ (1) period-based dependencies discovery
+ (2) time delay aggregation
+ This block can replace the self-attention family mechanism seamlessly.
+ """
+
+ def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False, configs=None):
+ super(AutoCorrelation, self).__init__()
+ print('Autocorrelation used !')
+ self.factor = factor
+ self.scale = scale
+ self.mask_flag = mask_flag
+ self.output_attention = output_attention
+ self.dropout = nn.Dropout(attention_dropout)
+ self.agg = None
+ self.use_wavelet = configs.wavelet
+
+ # @decor_time
+ def time_delay_agg_training(self, values, corr):
+ """
+ SpeedUp version of Autocorrelation (a batch-normalization style design)
+ This is for the training phase.
+ """
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+ index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
+ weights = torch.stack([mean_value[:, index[i]]
+ for i in range(top_k)], dim=-1)
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ pattern = torch.roll(tmp_values, -int(index[i]), -1)
+ delays_agg = delays_agg + pattern * \
+ (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(
+ 1).repeat(1, head, channel, length))
+ return delays_agg # size=[B, H, d, S]
+
+ def time_delay_agg_inference(self, values, corr):
+ """
+ SpeedUp version of Autocorrelation (a batch-normalization style design)
+ This is for the inference phase.
+ """
+ batch = values.shape[0]
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # index init
+ init_index = torch.arange(length).unsqueeze(0).unsqueeze(
+ 0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+ weights = torch.topk(mean_value, top_k, dim=-1)[0]
+ delay = torch.topk(mean_value, top_k, dim=-1)[1]
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values.repeat(1, 1, 1, 2)
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ tmp_delay = init_index + \
+ delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(
+ 1).repeat(1, head, channel, length)
+ pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+ delays_agg = delays_agg + pattern * \
+ (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(
+ 1).repeat(1, head, channel, length))
+ return delays_agg
+
+ def time_delay_agg_full(self, values, corr):
+ """
+ Standard version of Autocorrelation
+ """
+ batch = values.shape[0]
+ head = values.shape[1]
+ channel = values.shape[2]
+ length = values.shape[3]
+ # index init
+ init_index = torch.arange(length).unsqueeze(0).unsqueeze(
+ 0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
+ # find top k
+ top_k = int(self.factor * math.log(length))
+ weights = torch.topk(corr, top_k, dim=-1)[0]
+ delay = torch.topk(corr, top_k, dim=-1)[1]
+ # update corr
+ tmp_corr = torch.softmax(weights, dim=-1)
+ # aggregation
+ tmp_values = values.repeat(1, 1, 1, 2)
+ delays_agg = torch.zeros_like(values).float()
+ for i in range(top_k):
+ tmp_delay = init_index + delay[..., i].unsqueeze(-1)
+ pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+ delays_agg = delays_agg + pattern * \
+ (tmp_corr[..., i].unsqueeze(-1))
+ return delays_agg
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, H, E = queries.shape
+ _, S, _, D = values.shape
+ if L > S:
+ zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
+ values = torch.cat([values, zeros], dim=1)
+ keys = torch.cat([keys, zeros], dim=1)
+ else:
+ values = values[:, :L, :, :]
+ keys = keys[:, :L, :, :]
+
+ # period-based dependencies
+ if self.use_wavelet != 2:
+ if self.use_wavelet == 1:
+ j_list = self.j_list
+ queries = queries.reshape([B, L, -1])
+ keys = keys.reshape([B, L, -1])
+ Ql, Qh_list = self.dwt1d(
+ queries.transpose(1, 2)) # [B, H*D, L]
+ Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
+ qs = [queries.transpose(1, 2)] + Qh_list + [Ql] # [B, H*D, L]
+ ks = [keys.transpose(1, 2)] + Kh_list + [Kl]
+ q_list = []
+ k_list = []
+ for q, k, j in zip(qs, ks, j_list):
+ q_list += [interpolate(q, scale_factor=j,
+ mode='linear')[:, :, -L:]]
+ k_list += [interpolate(k, scale_factor=j,
+ mode='linear')[:, :, -L:]]
+ queries = torch.stack([i.reshape([B, H, E, L]) for i in q_list], dim=3).reshape(
+ [B, H, -1, L]).permute(0, 3, 1, 2)
+ keys = torch.stack([i.reshape([B, H, E, L]) for i in k_list], dim=3).reshape(
+ [B, H, -1, L]).permute(0, 3, 1, 2)
+ else:
+ pass
+ q_fft = torch.fft.rfft(queries.permute(
+ 0, 2, 3, 1).contiguous(), dim=-1) # size=[B, H, E, L]
+ k_fft = torch.fft.rfft(keys.permute(
+ 0, 2, 3, 1).contiguous(), dim=-1)
+ res = q_fft * torch.conj(k_fft)
+ corr = torch.fft.irfft(res, dim=-1) # size=[B, H, E, L]
+
+ # time delay agg
+ if self.training:
+ V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(
+ ), corr).permute(0, 3, 1, 2) # [B, L, H, E], [B, H, E, L] -> [B, L, H, E]
+ else:
+ V = self.time_delay_agg_inference(values.permute(
+ 0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+ else:
+ V_list = []
+ queries = queries.reshape([B, L, -1])
+ keys = keys.reshape([B, L, -1])
+ values = values.reshape([B, L, -1])
+ Ql, Qh_list = self.dwt1d(queries.transpose(1, 2)) # [B, H*D, L]
+ Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
+ Vl, Vh_list = self.dwt1d(values.transpose(1, 2))
+ qs = Qh_list + [Ql] # [B, H*D, L]
+ ks = Kh_list + [Kl]
+ vs = Vh_list + [Vl]
+ for q, k, v in zip(qs, ks, vs):
+ q = q.reshape([B, H, E, -1])
+ k = k.reshape([B, H, E, -1])
+ v = v.reshape([B, H, E, -1]).permute(0, 3, 1, 2)
+ q_fft = torch.fft.rfft(q.contiguous(), dim=-1)
+ k_fft = torch.fft.rfft(k.contiguous(), dim=-1)
+ res = q_fft * torch.conj(k_fft)
+ corr = torch.fft.irfft(res, dim=-1) # [B, H, E, L]
+ if self.training:
+ V = self.time_delay_agg_training(
+ v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+ else:
+ V = self.time_delay_agg_inference(
+ v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+ V_list += [V]
+ Vl = V_list[-1].reshape([B, -1, H*E]).transpose(1, 2)
+ Vh_list = [i.reshape([B, -1, H*E]).transpose(1, 2)
+ for i in V_list[:-1]]
+ V = self.dwt1div((Vl, Vh_list)).reshape(
+ [B, H, E, -1]).permute(0, 3, 1, 2)
+ # corr = self.dwt1div((V_list[-1], V_list[:-1]))
+
+ if self.output_attention:
+ # size = [B, L, H, E]
+ return (V.contiguous(), corr.permute(0, 3, 1, 2))
+ else:
+ return (V.contiguous(), None)
+
+
+class AutoCorrelationLayer(nn.Module):
+ def __init__(self, correlation, d_model, n_heads, d_keys=None,
+ d_values=None):
+ super(AutoCorrelationLayer, self).__init__()
+
+ d_keys = d_keys or (d_model // n_heads)
+ d_values = d_values or (d_model // n_heads)
+
+ self.inner_correlation = correlation
+ self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.value_projection = nn.Linear(d_model, d_values * n_heads)
+ self.out_projection = nn.Linear(d_values * n_heads, d_model)
+ self.n_heads = n_heads
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, _ = queries.shape
+ _, S, _ = keys.shape
+ H = self.n_heads
+
+ queries = self.query_projection(queries).view(B, L, H, -1)
+ keys = self.key_projection(keys).view(B, S, H, -1)
+ values = self.value_projection(values).view(B, S, H, -1)
+
+ out, attn = self.inner_correlation(
+ queries,
+ keys,
+ values,
+ attn_mask
+ )
+
+ out = out.view(B, L, -1)
+ return self.out_projection(out), attn
diff --git a/basicts/archs/arch_zoo/fedformer_arch/embed.py b/basicts/archs/arch_zoo/fedformer_arch/embed.py
new file mode 100644
index 0000000..b08f259
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/embed.py
@@ -0,0 +1,115 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+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', bias=False)
+ 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, num_time_features):
+ super(TimeFeatureEmbedding, self).__init__()
+ # freq_map = {'h': 2}
+ # NOTE: this is different from official autoformer, since METR-LA, PEMS-BAY, and PEMS0X datasets only contain two kind of temporal features (time in day, day in week)
+ self.embed = nn.Linear(num_time_features, d_model, bias=False)
+
+ def forward(self, x):
+ return self.embed(x)
+
+
+class DataEmbedding(nn.Module):
+ def __init__(self, c_in, d_model, num_time_features, 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=d_model, num_time_features=num_time_features)
+ self.dropout = nn.Dropout(p=dropout)
+
+ def forward(self, x, x_mark):
+ x = self.value_embedding(
+ x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
+ return self.dropout(x)
+
+
+class DataEmbedding_onlypos(nn.Module):
+ def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+ super(DataEmbedding_onlypos, self).__init__()
+
+ self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+ self.position_embedding = PositionalEmbedding(d_model=d_model)
+ self.dropout = nn.Dropout(p=dropout)
+
+ def forward(self, x, x_mark):
+ x = self.value_embedding(x) + self.position_embedding(x)
+ return self.dropout(x)
+
+
+class DataEmbedding_wo_pos(nn.Module):
+ def __init__(self, c_in, d_model, num_time_features, dropout=0.1):
+ super(DataEmbedding_wo_pos, 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=d_model, num_time_features=num_time_features)
+ self.dropout = nn.Dropout(p=dropout)
+
+ def forward(self, x, x_mark):
+ x = self.value_embedding(x) + self.temporal_embedding(x_mark)
+ return self.dropout(x)
diff --git a/basicts/archs/arch_zoo/fedformer_arch/fedformer_arch.py b/basicts/archs/arch_zoo/fedformer_arch/fedformer_arch.py
new file mode 100644
index 0000000..c124a3e
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/fedformer_arch.py
@@ -0,0 +1,183 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .embed import DataEmbedding_wo_pos, DataEmbedding
+from .auto_correlation import AutoCorrelationLayer
+from .fourier_correlation import FourierBlock, FourierCrossAttention
+from .multi_wavelet_correlation import MultiWaveletCross, MultiWaveletTransform
+from .fedformer_enc_dec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp, series_decomp_multi
+from ..utils import data_transformation_4_xformer
+
+
+class FEDformer(nn.Module):
+ """
+ Paper: FEDformer: Frequency Enhanced Decomposed Transformer for Long-term Series Forecasting
+ Link: https://arxiv.org/abs/2201.12740v3
+ Ref Official Code: https://github.com/MAZiqing/FEDformer
+ """
+
+ def __init__(self, **model_args):
+ super(FEDformer, self).__init__()
+ self.version = model_args["version"]
+ self.mode_select = model_args["mode_select"]
+ self.modes = model_args["modes"]
+ self.seq_len = int(model_args["seq_len"])
+ self.label_len = int(model_args["label_len"])
+ self.pred_len = int(model_args["pred_len"])
+ self.output_attention = model_args["output_attention"]
+ self.embedding_type = model_args["embedding_type"]
+
+ # Decomp
+ kernel_size = model_args["moving_avg"]
+ self.decomp = series_decomp(kernel_size)
+
+ # Decomp
+ kernel_size = model_args["moving_avg"]
+ if isinstance(kernel_size, list):
+ self.decomp = series_decomp_multi(kernel_size)
+ else:
+ self.decomp = series_decomp(kernel_size)
+
+ # Embedding
+ # The series-wise connection inherently contains the sequential information.
+ # Thus, we can discard the position embedding of transformers.
+ if self.embedding_type == "DataEmbedding_wo_pos":
+ self.enc_embedding = DataEmbedding_wo_pos(model_args["enc_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ self.dec_embedding = DataEmbedding_wo_pos(model_args["dec_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ elif self.embedding_type == "DataEmbedding":
+ self.enc_embedding = DataEmbedding(model_args["enc_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ self.dec_embedding = DataEmbedding(model_args["dec_in"], model_args["d_model"], model_args["num_time_features"], model_args["dropout"])
+ else:
+ raise Exception("Unknown embedding type.")
+
+ if model_args["version"] == 'Wavelets':
+ encoder_self_att = MultiWaveletTransform(
+ ich=model_args["d_model"], L=model_args["L"], base=model_args["base"])
+ decoder_self_att = MultiWaveletTransform(
+ ich=model_args["d_model"], L=model_args["L"], base=model_args["base"])
+ decoder_cross_att = MultiWaveletCross(in_channels=model_args["d_model"],
+ out_channels=model_args["d_model"],
+ seq_len_q=self.seq_len // 2 + self.pred_len,
+ seq_len_kv=self.seq_len,
+ modes=model_args["modes"],
+ ich=model_args["d_model"],
+ base=model_args["base"],
+ activation=model_args["cross_activation"])
+ else:
+ encoder_self_att = FourierBlock(in_channels=model_args["d_model"],
+ out_channels=model_args["d_model"],
+ seq_len=self.seq_len,
+ modes=model_args["modes"],
+ mode_select_method=model_args["mode_select"])
+ decoder_self_att = FourierBlock(in_channels=model_args["d_model"],
+ out_channels=model_args["d_model"],
+ seq_len=self.seq_len//2+self.pred_len,
+ modes=model_args["modes"],
+ mode_select_method=model_args["mode_select"])
+ decoder_cross_att = FourierCrossAttention(in_channels=model_args["d_model"],
+ out_channels=model_args["d_model"],
+ seq_len_q=self.seq_len//2+self.pred_len,
+ seq_len_kv=self.seq_len,
+ modes=model_args["modes"],
+ mode_select_method=model_args["mode_select"])
+ # Encoder
+ enc_modes = int(min(model_args["modes"], model_args["seq_len"]//2))
+ dec_modes = int(
+ min(model_args["modes"], (model_args["seq_len"]//2+model_args["pred_len"])//2))
+ print('enc_modes: {}, dec_modes: {}'.format(enc_modes, dec_modes))
+
+ self.encoder = Encoder(
+ [
+ EncoderLayer(
+ AutoCorrelationLayer(
+ encoder_self_att,
+ model_args["d_model"], model_args["n_heads"]),
+
+ model_args["d_model"],
+ model_args["d_ff"],
+ moving_avg=model_args["moving_avg"],
+ dropout=model_args["dropout"],
+ activation=model_args["activation"]
+ ) for l in range(model_args["e_layers"])
+ ],
+ norm_layer=my_Layernorm(model_args["d_model"])
+ )
+ # Decoder
+ self.decoder = Decoder(
+ [
+ DecoderLayer(
+ AutoCorrelationLayer(
+ decoder_self_att,
+ model_args["d_model"], model_args["n_heads"]),
+ AutoCorrelationLayer(
+ decoder_cross_att,
+ model_args["d_model"], model_args["n_heads"]),
+ model_args["d_model"],
+ model_args["c_out"],
+ model_args["d_ff"],
+ moving_avg=model_args["moving_avg"],
+ dropout=model_args["dropout"],
+ activation=model_args["activation"],
+ )
+ for l in range(model_args["d_layers"])
+ ],
+ norm_layer=my_Layernorm(model_args["d_model"]),
+ projection=nn.Linear(model_args["d_model"], model_args["c_out"], bias=True)
+ )
+
+ def forward_xformer(self, x_enc: torch.Tensor, x_mark_enc: torch.Tensor, x_dec: torch.Tensor, x_mark_dec: torch.Tensor,
+ enc_self_mask: torch.Tensor=None, dec_self_mask: torch.Tensor=None, dec_enc_mask: torch.Tensor=None) -> torch.Tensor:
+ """Feed forward of FEDformer. Kindly note that `enc_self_mask`, `dec_self_mask`, and `dec_enc_mask` are not actually used in FEDFormer.
+
+ Args:
+ x_enc (torch.Tensor): input data of encoder (without the time features). Shape: [B, L1, N]
+ x_mark_enc (torch.Tensor): time features input of encoder w.r.t. x_enc. Shape: [B, L1, C-1]
+ x_dec (torch.Tensor): input data of decoder. Shape: [B, start_token_length + L2, N]
+ x_mark_dec (torch.Tensor): time features input to decoder w.r.t. x_dec. Shape: [B, start_token_length + L2, C-1]
+ enc_self_mask (torch.Tensor, optional): encoder self attention masks. Defaults to None.
+ dec_self_mask (torch.Tensor, optional): decoder self attention masks. Defaults to None.
+ dec_enc_mask (torch.Tensor, optional): decoder encoder self attention masks. Defaults to None.
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L2, N, 1]
+ """
+
+ # decomp init
+ mean = torch.mean(x_enc, dim=1).unsqueeze(
+ 1).repeat(1, self.pred_len, 1)
+ zeros = torch.zeros([x_dec.shape[0], self.pred_len,
+ x_dec.shape[2]]).to(mean.device)
+ seasonal_init, trend_init = self.decomp(x_enc)
+ # decoder input
+ trend_init = torch.cat(
+ [trend_init[:, -self.label_len:, :], mean], dim=1)
+ seasonal_init = F.pad(
+ seasonal_init[:, -self.label_len:, :], (0, 0, 0, self.pred_len))
+ # enc
+ enc_out = self.enc_embedding(x_enc, x_mark_enc)
+ enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+ # dec
+ dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
+ seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
+ trend=trend_init)
+ # final
+ dec_out = trend_part + seasonal_part
+
+ return dec_out[:, -self.pred_len:, :].unsqueeze(-1) # [B, L, N, C]
+
+ def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool, **kwargs) -> torch.Tensor:
+ """
+
+ Args:
+ history_data (Tensor): Input data with shape: [B, L1, N, C]
+ future_data (Tensor): Future data with shape: [B, L2, N, C]
+
+ Returns:
+ torch.Tensor: outputs with shape [B, L2, N, 1]
+ """
+
+ x_enc, x_mark_enc, x_dec, x_mark_dec = data_transformation_4_xformer(history_data=history_data, future_data=future_data, start_token_len=self.label_len)
+ prediction = self.forward_xformer(x_enc=x_enc, x_mark_enc=x_mark_enc, x_dec=x_dec, x_mark_dec=x_mark_dec)
+ return prediction
+
diff --git a/basicts/archs/arch_zoo/fedformer_arch/fedformer_enc_dec.py b/basicts/archs/arch_zoo/fedformer_arch/fedformer_enc_dec.py
new file mode 100644
index 0000000..13b5921
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/fedformer_enc_dec.py
@@ -0,0 +1,219 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class my_Layernorm(nn.Module):
+ """
+ Special designed layernorm for the seasonal part
+ """
+ def __init__(self, channels):
+ super(my_Layernorm, self).__init__()
+ self.layernorm = nn.LayerNorm(channels)
+
+ def forward(self, x):
+ x_hat = self.layernorm(x)
+ bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
+ return x_hat - bias
+
+
+class moving_avg(nn.Module):
+ """
+ Moving average block to highlight the trend of time series
+ """
+ def __init__(self, kernel_size, stride):
+ super(moving_avg, self).__init__()
+ self.kernel_size = kernel_size
+ self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+ def forward(self, x):
+ # padding on the both ends of time series
+ front = x[:, 0:1, :].repeat(1, self.kernel_size - 1-math.floor((self.kernel_size - 1) // 2), 1)
+ end = x[:, -1:, :].repeat(1, math.floor((self.kernel_size - 1) // 2), 1)
+ x = torch.cat([front, x, end], dim=1)
+ x = self.avg(x.permute(0, 2, 1))
+ x = x.permute(0, 2, 1)
+ return x
+
+
+class series_decomp(nn.Module):
+ """
+ Series decomposition block
+ """
+ def __init__(self, kernel_size):
+ super(series_decomp, self).__init__()
+ self.moving_avg = moving_avg(kernel_size, stride=1)
+
+ def forward(self, x):
+ moving_mean = self.moving_avg(x)
+ res = x - moving_mean
+ return res, moving_mean
+
+
+class series_decomp_multi(nn.Module):
+ """
+ Series decomposition block
+ """
+ def __init__(self, kernel_size):
+ super(series_decomp_multi, self).__init__()
+ self.moving_avg = [moving_avg(kernel, stride=1) for kernel in kernel_size]
+ self.layer = torch.nn.Linear(1, len(kernel_size))
+
+ def forward(self, x):
+ moving_mean=[]
+ for func in self.moving_avg:
+ moving_avg = func(x)
+ moving_mean.append(moving_avg.unsqueeze(-1))
+ moving_mean=torch.cat(moving_mean,dim=-1)
+ moving_mean = torch.sum(moving_mean*nn.Softmax(-1)(self.layer(x.unsqueeze(-1))),dim=-1)
+ res = x - moving_mean
+ return res, moving_mean
+
+
+class FourierDecomp(nn.Module):
+ def __init__(self):
+ super(FourierDecomp, self).__init__()
+ pass
+
+ def forward(self, x):
+ x_ft = torch.fft.rfft(x, dim=-1)
+
+
+class EncoderLayer(nn.Module):
+ """
+ Autoformer encoder layer with the progressive decomposition architecture
+ """
+ def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
+ super(EncoderLayer, self).__init__()
+ d_ff = d_ff or 4 * d_model
+ self.attention = attention
+ self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+ self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+
+ if isinstance(moving_avg, list):
+ self.decomp1 = series_decomp_multi(moving_avg)
+ self.decomp2 = series_decomp_multi(moving_avg)
+ else:
+ self.decomp1 = series_decomp(moving_avg)
+ self.decomp2 = series_decomp(moving_avg)
+
+ self.dropout = nn.Dropout(dropout)
+ self.activation = F.relu if activation == "relu" else F.gelu
+
+ def forward(self, x, attn_mask=None):
+ new_x, attn = self.attention(
+ x, x, x,
+ attn_mask=attn_mask
+ )
+ x = x + self.dropout(new_x)
+ x, _ = self.decomp1(x)
+ y = x
+ y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+ y = self.dropout(self.conv2(y).transpose(-1, 1))
+ res, _ = self.decomp2(x + y)
+ return res, attn
+
+
+class Encoder(nn.Module):
+ """
+ Autoformer encoder
+ """
+ def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+ super(Encoder, self).__init__()
+ self.attn_layers = nn.ModuleList(attn_layers)
+ self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+ self.norm = norm_layer
+
+ def forward(self, x, attn_mask=None):
+ attns = []
+ if self.conv_layers is not None:
+ for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
+ x, attn = attn_layer(x, attn_mask=attn_mask)
+ x = conv_layer(x)
+ attns.append(attn)
+ x, attn = self.attn_layers[-1](x)
+ attns.append(attn)
+ else:
+ for attn_layer in self.attn_layers:
+ x, attn = attn_layer(x, attn_mask=attn_mask)
+ attns.append(attn)
+
+ if self.norm is not None:
+ x = self.norm(x)
+
+ return x, attns
+
+
+class DecoderLayer(nn.Module):
+ """
+ Autoformer decoder layer with the progressive decomposition architecture
+ """
+ def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
+ moving_avg=25, dropout=0.1, activation="relu"):
+ super(DecoderLayer, self).__init__()
+ d_ff = d_ff or 4 * d_model
+ self.self_attention = self_attention
+ self.cross_attention = cross_attention
+ self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+ self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+
+ if isinstance(moving_avg, list):
+ self.decomp1 = series_decomp_multi(moving_avg)
+ self.decomp2 = series_decomp_multi(moving_avg)
+ self.decomp3 = series_decomp_multi(moving_avg)
+ else:
+ self.decomp1 = series_decomp(moving_avg)
+ self.decomp2 = series_decomp(moving_avg)
+ self.decomp3 = series_decomp(moving_avg)
+
+ self.dropout = nn.Dropout(dropout)
+ self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
+ padding_mode='circular', bias=False)
+ self.activation = F.relu if activation == "relu" else F.gelu
+
+ def forward(self, x, cross, x_mask=None, cross_mask=None):
+ x = x + self.dropout(self.self_attention(
+ x, x, x,
+ attn_mask=x_mask
+ )[0])
+
+ x, trend1 = self.decomp1(x)
+ x = x + self.dropout(self.cross_attention(
+ x, cross, cross,
+ attn_mask=cross_mask
+ )[0])
+
+ x, trend2 = self.decomp2(x)
+ y = x
+ y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+ y = self.dropout(self.conv2(y).transpose(-1, 1))
+ x, trend3 = self.decomp3(x + y)
+
+ residual_trend = trend1 + trend2 + trend3
+ residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
+ return x, residual_trend
+
+
+class Decoder(nn.Module):
+ """
+ Autoformer encoder
+ """
+ def __init__(self, layers, norm_layer=None, projection=None):
+ super(Decoder, self).__init__()
+ self.layers = nn.ModuleList(layers)
+ self.norm = norm_layer
+ self.projection = projection
+
+ def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
+ for layer in self.layers:
+ x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+ trend = trend + residual_trend
+
+ if self.norm is not None:
+ x = self.norm(x)
+
+ if self.projection is not None:
+ x = self.projection(x)
+ return x, trend
diff --git a/basicts/archs/arch_zoo/fedformer_arch/fourier_correlation.py b/basicts/archs/arch_zoo/fedformer_arch/fourier_correlation.py
new file mode 100644
index 0000000..5e19d97
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/fourier_correlation.py
@@ -0,0 +1,136 @@
+# coding=utf-8
+# author=maziqing
+# email=maziqing.mzq@alibaba-inc.com
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+def get_frequency_modes(seq_len, modes=64, mode_select_method='random'):
+ """
+ get modes on frequency domain:
+ 'random' means sampling randomly;
+ 'else' means sampling the lowest modes;
+ """
+ modes = min(modes, seq_len//2)
+ if mode_select_method == 'random':
+ index = list(range(0, seq_len // 2))
+ np.random.shuffle(index)
+ index = index[:modes]
+ else:
+ index = list(range(0, modes))
+ index.sort()
+ return index
+
+
+# ########## fourier layer #############
+class FourierBlock(nn.Module):
+ def __init__(self, in_channels, out_channels, seq_len, modes=0, mode_select_method='random'):
+ super(FourierBlock, self).__init__()
+ print('fourier enhanced block used!')
+ """
+ 1D Fourier block. It performs representation learning on frequency domain,
+ it does FFT, linear transform, and Inverse FFT.
+ """
+ # get modes on frequency domain
+ self.index = get_frequency_modes(
+ seq_len, modes=modes, mode_select_method=mode_select_method)
+ print('modes={}, index={}'.format(modes, self.index))
+
+ self.scale = (1 / (in_channels * out_channels))
+ self.weights1 = nn.Parameter(
+ self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index), dtype=torch.cfloat))
+
+ # Complex multiplication
+ def compl_mul1d(self, input, weights):
+ # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+ return torch.einsum("bhi,hio->bho", input, weights)
+
+ def forward(self, q, k, v, mask):
+ # size = [B, L, H, E]
+ B, L, H, E = q.shape
+ x = q.permute(0, 2, 3, 1)
+ # Compute Fourier coefficients
+ x_ft = torch.fft.rfft(x, dim=-1)
+ # Perform Fourier neural operations
+ out_ft = torch.zeros(B, H, E, L // 2 + 1,
+ device=x.device, dtype=torch.cfloat)
+ for wi, i in enumerate(self.index):
+ out_ft[:, :, :, wi] = self.compl_mul1d(
+ x_ft[:, :, :, i], self.weights1[:, :, :, wi])
+ # Return to time domain
+ x = torch.fft.irfft(out_ft, n=x.size(-1))
+ return (x, None)
+
+
+# ########## Fourier Cross Former ####################
+class FourierCrossAttention(nn.Module):
+ def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=64, mode_select_method='random',
+ activation='tanh', policy=0):
+ super(FourierCrossAttention, self).__init__()
+ print(' fourier enhanced cross attention used!')
+ """
+ 1D Fourier Cross Attention layer. It does FFT, linear transform, attention mechanism and Inverse FFT.
+ """
+ self.activation = activation
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ # get modes for queries and keys (& values) on frequency domain
+ self.index_q = get_frequency_modes(
+ seq_len_q, modes=modes, mode_select_method=mode_select_method)
+ self.index_kv = get_frequency_modes(
+ seq_len_kv, modes=modes, mode_select_method=mode_select_method)
+
+ print('modes_q={}, index_q={}'.format(len(self.index_q), self.index_q))
+ print('modes_kv={}, index_kv={}'.format(
+ len(self.index_kv), self.index_kv))
+
+ self.scale = (1 / (in_channels * out_channels))
+ self.weights1 = nn.Parameter(
+ self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index_q), dtype=torch.cfloat))
+
+ # Complex multiplication
+ def compl_mul1d(self, input, weights):
+ # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+ return torch.einsum("bhi,hio->bho", input, weights)
+
+ def forward(self, q, k, v, mask):
+ # size = [B, L, H, E]
+ B, L, H, E = q.shape
+ xq = q.permute(0, 2, 3, 1) # size = [B, H, E, L]
+ xk = k.permute(0, 2, 3, 1)
+ xv = v.permute(0, 2, 3, 1)
+
+ # Compute Fourier coefficients
+ xq_ft_ = torch.zeros(B, H, E, len(self.index_q),
+ device=xq.device, dtype=torch.cfloat)
+ xq_ft = torch.fft.rfft(xq, dim=-1)
+ for i, j in enumerate(self.index_q):
+ xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
+ xk_ft_ = torch.zeros(B, H, E, len(self.index_kv),
+ device=xq.device, dtype=torch.cfloat)
+ xk_ft = torch.fft.rfft(xk, dim=-1)
+ for i, j in enumerate(self.index_kv):
+ xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
+
+ # perform attention mechanism on frequency domain
+ xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
+ if self.activation == 'tanh':
+ xqk_ft = xqk_ft.tanh()
+ elif self.activation == 'softmax':
+ xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
+ xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
+ else:
+ raise Exception(
+ '{} actiation function is not implemented'.format(self.activation))
+ xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
+ xqkvw = torch.einsum("bhex,heox->bhox", xqkv_ft, self.weights1)
+ out_ft = torch.zeros(B, H, E, L // 2 + 1,
+ device=xq.device, dtype=torch.cfloat)
+ for i, j in enumerate(self.index_q):
+ out_ft[:, :, :, j] = xqkvw[:, :, :, i]
+ # Return to time domain
+ out = torch.fft.irfft(out_ft / self.in_channels /
+ self.out_channels, n=xq.size(-1))
+ return (out, None)
diff --git a/basicts/archs/arch_zoo/fedformer_arch/masking.py b/basicts/archs/arch_zoo/fedformer_arch/masking.py
new file mode 100644
index 0000000..49b4efc
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/masking.py
@@ -0,0 +1,47 @@
+import math
+
+import torch
+import numpy as np
+
+
+class TriangularCausalMask():
+ def __init__(self, B, L, device="cpu"):
+ mask_shape = [B, 1, L, L]
+ with torch.no_grad():
+ self._mask = torch.triu(torch.ones(
+ mask_shape, dtype=torch.bool), diagonal=1).to(device)
+
+ @property
+ def mask(self):
+ return self._mask
+
+
+class ProbMask():
+ def __init__(self, B, H, L, index, scores, device="cpu"):
+ _mask = torch.ones(
+ L, scores.shape[-1], dtype=torch.bool).to(device).triu(1)
+ _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])
+ indicator = _mask_ex[torch.arange(B)[:, None, None],
+ torch.arange(H)[None, :, None],
+ index, :].to(device)
+ self._mask = indicator.view(scores.shape).to(device)
+
+ @property
+ def mask(self):
+ return self._mask
+
+
+class LocalMask():
+ def __init__(self, B, L, S, device="cpu"):
+ mask_shape = [B, 1, L, S]
+ with torch.no_grad():
+ self.len = math.ceil(np.log2(L))
+ self._mask1 = torch.triu(torch.ones(
+ mask_shape, dtype=torch.bool), diagonal=1).to(device)
+ self._mask2 = ~torch.triu(torch.ones(
+ mask_shape, dtype=torch.bool), diagonal=-self.len).to(device)
+ self._mask = self._mask1+self._mask2
+
+ @property
+ def mask(self):
+ return self._mask
diff --git a/basicts/archs/arch_zoo/fedformer_arch/multi_wavelet_correlation.py b/basicts/archs/arch_zoo/fedformer_arch/multi_wavelet_correlation.py
new file mode 100644
index 0000000..7e1e51f
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/multi_wavelet_correlation.py
@@ -0,0 +1,385 @@
+import math
+from typing import List, Tuple
+
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from .utils import get_filter
+
+
+class MultiWaveletTransform(nn.Module):
+ """
+ 1D multiwavelet block.
+ """
+
+ def __init__(self, ich=1, k=8, alpha=16, c=128,
+ nCZ=1, L=0, base='legendre', attention_dropout=0.1):
+ super(MultiWaveletTransform, self).__init__()
+ print('base', base)
+ self.k = k
+ self.c = c
+ self.L = L
+ self.nCZ = nCZ
+ self.Lk0 = nn.Linear(ich, c * k)
+ self.Lk1 = nn.Linear(c * k, ich)
+ self.ich = ich
+ self.MWT_CZ = nn.ModuleList(
+ MWT_CZ1d(k, alpha, L, c, base) for i in range(nCZ))
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, H, E = queries.shape
+ _, S, _, D = values.shape
+ if L > S:
+ zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
+ values = torch.cat([values, zeros], dim=1)
+ keys = torch.cat([keys, zeros], dim=1)
+ else:
+ values = values[:, :L, :, :]
+ keys = keys[:, :L, :, :]
+ values = values.view(B, L, -1)
+
+ V = self.Lk0(values).view(B, L, self.c, -1)
+ for i in range(self.nCZ):
+ V = self.MWT_CZ[i](V)
+ if i < self.nCZ - 1:
+ V = F.relu(V)
+
+ V = self.Lk1(V.view(B, L, -1))
+ V = V.view(B, L, -1, D)
+ return (V.contiguous(), None)
+
+
+class MultiWaveletCross(nn.Module):
+ """
+ 1D Multiwavelet Cross Attention layer.
+ """
+
+ def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes, c=64,
+ k=8, ich=512,
+ L=0,
+ base='legendre',
+ mode_select_method='random',
+ initializer=None, activation='tanh',
+ **kwargs):
+ super(MultiWaveletCross, self).__init__()
+ print('base', base)
+
+ self.c = c
+ self.k = k
+ self.L = L
+ H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
+ H0r = H0 @ PHI0
+ G0r = G0 @ PHI0
+ H1r = H1 @ PHI1
+ G1r = G1 @ PHI1
+
+ H0r[np.abs(H0r) < 1e-8] = 0
+ H1r[np.abs(H1r) < 1e-8] = 0
+ G0r[np.abs(G0r) < 1e-8] = 0
+ G1r[np.abs(G1r) < 1e-8] = 0
+ self.max_item = 3
+
+ self.attn1 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
+ seq_len_kv=seq_len_kv, modes=modes, activation=activation,
+ mode_select_method=mode_select_method)
+ self.attn2 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
+ seq_len_kv=seq_len_kv, modes=modes, activation=activation,
+ mode_select_method=mode_select_method)
+ self.attn3 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
+ seq_len_kv=seq_len_kv, modes=modes, activation=activation,
+ mode_select_method=mode_select_method)
+ self.attn4 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
+ seq_len_kv=seq_len_kv, modes=modes, activation=activation,
+ mode_select_method=mode_select_method)
+ self.T0 = nn.Linear(k, k)
+ self.register_buffer('ec_s', torch.Tensor(
+ np.concatenate((H0.T, H1.T), axis=0)))
+ self.register_buffer('ec_d', torch.Tensor(
+ np.concatenate((G0.T, G1.T), axis=0)))
+
+ self.register_buffer('rc_e', torch.Tensor(
+ np.concatenate((H0r, G0r), axis=0)))
+ self.register_buffer('rc_o', torch.Tensor(
+ np.concatenate((H1r, G1r), axis=0)))
+
+ self.Lk = nn.Linear(ich, c * k)
+ self.Lq = nn.Linear(ich, c * k)
+ self.Lv = nn.Linear(ich, c * k)
+ self.out = nn.Linear(c * k, ich)
+ self.modes1 = modes
+
+ def forward(self, q, k, v, mask=None):
+ B, N, H, E = q.shape # (B, N, H, E) torch.Size([3, 768, 8, 2])
+ _, S, _, _ = k.shape # (B, S, H, E) torch.Size([3, 96, 8, 2])
+
+ q = q.view(q.shape[0], q.shape[1], -1)
+ k = k.view(k.shape[0], k.shape[1], -1)
+ v = v.view(v.shape[0], v.shape[1], -1)
+ q = self.Lq(q)
+ q = q.view(q.shape[0], q.shape[1], self.c, self.k)
+ k = self.Lk(k)
+ k = k.view(k.shape[0], k.shape[1], self.c, self.k)
+ v = self.Lv(v)
+ v = v.view(v.shape[0], v.shape[1], self.c, self.k)
+
+ if N > S:
+ zeros = torch.zeros_like(q[:, :(N - S), :]).float()
+ v = torch.cat([v, zeros], dim=1)
+ k = torch.cat([k, zeros], dim=1)
+ else:
+ v = v[:, :N, :, :]
+ k = k[:, :N, :, :]
+
+ ns = math.floor(np.log2(N))
+ nl = pow(2, math.ceil(np.log2(N)))
+ extra_q = q[:, 0:nl - N, :, :]
+ extra_k = k[:, 0:nl - N, :, :]
+ extra_v = v[:, 0:nl - N, :, :]
+ q = torch.cat([q, extra_q], 1)
+ k = torch.cat([k, extra_k], 1)
+ v = torch.cat([v, extra_v], 1)
+
+ Ud_q = torch.jit.annotate(List[Tuple[Tensor]], [])
+ Ud_k = torch.jit.annotate(List[Tuple[Tensor]], [])
+ Ud_v = torch.jit.annotate(List[Tuple[Tensor]], [])
+
+ Us_q = torch.jit.annotate(List[Tensor], [])
+ Us_k = torch.jit.annotate(List[Tensor], [])
+ Us_v = torch.jit.annotate(List[Tensor], [])
+
+ Ud = torch.jit.annotate(List[Tensor], [])
+ Us = torch.jit.annotate(List[Tensor], [])
+
+ # decompose
+ for i in range(ns - self.L):
+ # print('q shape',q.shape)
+ d, q = self.wavelet_transform(q)
+ Ud_q += [tuple([d, q])]
+ Us_q += [d]
+ for i in range(ns - self.L):
+ d, k = self.wavelet_transform(k)
+ Ud_k += [tuple([d, k])]
+ Us_k += [d]
+ for i in range(ns - self.L):
+ d, v = self.wavelet_transform(v)
+ Ud_v += [tuple([d, v])]
+ Us_v += [d]
+ for i in range(ns - self.L):
+ dk, sk = Ud_k[i], Us_k[i]
+ dq, sq = Ud_q[i], Us_q[i]
+ dv, sv = Ud_v[i], Us_v[i]
+ Ud += [self.attn1(dq[0], dk[0], dv[0], mask)[0] +
+ self.attn2(dq[1], dk[1], dv[1], mask)[0]]
+ Us += [self.attn3(sq, sk, sv, mask)[0]]
+ v = self.attn4(q, k, v, mask)[0]
+
+ # reconstruct
+ for i in range(ns - 1 - self.L, -1, -1):
+ v = v + Us[i]
+ v = torch.cat((v, Ud[i]), -1)
+ v = self.evenOdd(v)
+ v = self.out(v[:, :N, :, :].contiguous().view(B, N, -1))
+ return (v.contiguous(), None)
+
+ def wavelet_transform(self, x):
+ xa = torch.cat([x[:, ::2, :, :],
+ x[:, 1::2, :, :],
+ ], -1)
+ d = torch.matmul(xa, self.ec_d)
+ s = torch.matmul(xa, self.ec_s)
+ return d, s
+
+ def evenOdd(self, x):
+ B, N, c, ich = x.shape # (B, N, c, k)
+ assert ich == 2 * self.k
+ x_e = torch.matmul(x, self.rc_e)
+ x_o = torch.matmul(x, self.rc_o)
+
+ x = torch.zeros(B, N * 2, c, self.k,
+ device=x.device)
+ x[..., ::2, :, :] = x_e
+ x[..., 1::2, :, :] = x_o
+ return x
+
+
+class FourierCrossAttentionW(nn.Module):
+ def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=16, activation='tanh',
+ mode_select_method='random'):
+ super(FourierCrossAttentionW, self).__init__()
+ print('corss fourier correlation used!')
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.modes1 = modes
+ self.activation = activation
+
+ def forward(self, q, k, v, mask):
+ B, L, E, H = q.shape
+
+ # size = [B, H, E, L] torch.Size([3, 8, 64, 512])
+ xq = q.permute(0, 3, 2, 1)
+ xk = k.permute(0, 3, 2, 1)
+ xv = v.permute(0, 3, 2, 1)
+ self.index_q = list(range(0, min(int(L // 2), self.modes1)))
+ self.index_k_v = list(
+ range(0, min(int(xv.shape[3] // 2), self.modes1)))
+
+ # Compute Fourier coefficients
+ xq_ft_ = torch.zeros(B, H, E, len(self.index_q),
+ device=xq.device, dtype=torch.cfloat)
+ xq_ft = torch.fft.rfft(xq, dim=-1)
+ for i, j in enumerate(self.index_q):
+ xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
+
+ xk_ft_ = torch.zeros(B, H, E, len(self.index_k_v),
+ device=xq.device, dtype=torch.cfloat)
+ xk_ft = torch.fft.rfft(xk, dim=-1)
+ for i, j in enumerate(self.index_k_v):
+ xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
+ xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
+ if self.activation == 'tanh':
+ xqk_ft = xqk_ft.tanh()
+ elif self.activation == 'softmax':
+ xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
+ xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
+ else:
+ raise Exception(
+ '{} actiation function is not implemented'.format(self.activation))
+ xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
+
+ xqkvw = xqkv_ft
+ out_ft = torch.zeros(B, H, E, L // 2 + 1,
+ device=xq.device, dtype=torch.cfloat)
+ for i, j in enumerate(self.index_q):
+ out_ft[:, :, :, j] = xqkvw[:, :, :, i]
+
+ out = torch.fft.irfft(out_ft / self.in_channels /
+ self.out_channels, n=xq.size(-1)).permute(0, 3, 2, 1)
+ # size = [B, L, H, E]
+ return (out, None)
+
+
+class sparseKernelFT1d(nn.Module):
+ def __init__(self,
+ k, alpha, c=1,
+ nl=1,
+ initializer=None,
+ **kwargs):
+ super(sparseKernelFT1d, self).__init__()
+
+ self.modes1 = alpha
+ self.scale = (1 / (c * k * c * k))
+ self.weights1 = nn.Parameter(
+ self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.cfloat))
+ self.weights1.requires_grad = True
+ self.k = k
+
+ def compl_mul1d(self, x, weights):
+ # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+ return torch.einsum("bix,iox->box", x, weights)
+
+ def forward(self, x):
+ B, N, c, k = x.shape # (B, N, c, k)
+
+ x = x.view(B, N, -1)
+ x = x.permute(0, 2, 1)
+ x_fft = torch.fft.rfft(x)
+ # Multiply relevant Fourier modes
+ l = min(self.modes1, N // 2 + 1)
+ # l = N//2+1
+ out_ft = torch.zeros(B, c * k, N // 2 + 1,
+ device=x.device, dtype=torch.cfloat)
+ out_ft[:, :, :l] = self.compl_mul1d(
+ x_fft[:, :, :l], self.weights1[:, :, :l])
+ x = torch.fft.irfft(out_ft, n=N)
+ x = x.permute(0, 2, 1).view(B, N, c, k)
+ return x
+
+
+# ##
+class MWT_CZ1d(nn.Module):
+ def __init__(self,
+ k=3, alpha=64,
+ L=0, c=1,
+ base='legendre',
+ initializer=None,
+ **kwargs):
+ super(MWT_CZ1d, self).__init__()
+
+ self.k = k
+ self.L = L
+ H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
+ H0r = H0 @ PHI0
+ G0r = G0 @ PHI0
+ H1r = H1 @ PHI1
+ G1r = G1 @ PHI1
+
+ H0r[np.abs(H0r) < 1e-8] = 0
+ H1r[np.abs(H1r) < 1e-8] = 0
+ G0r[np.abs(G0r) < 1e-8] = 0
+ G1r[np.abs(G1r) < 1e-8] = 0
+ self.max_item = 3
+
+ self.A = sparseKernelFT1d(k, alpha, c)
+ self.B = sparseKernelFT1d(k, alpha, c)
+ self.C = sparseKernelFT1d(k, alpha, c)
+
+ self.T0 = nn.Linear(k, k)
+
+ self.register_buffer('ec_s', torch.Tensor(
+ np.concatenate((H0.T, H1.T), axis=0)))
+ self.register_buffer('ec_d', torch.Tensor(
+ np.concatenate((G0.T, G1.T), axis=0)))
+
+ self.register_buffer('rc_e', torch.Tensor(
+ np.concatenate((H0r, G0r), axis=0)))
+ self.register_buffer('rc_o', torch.Tensor(
+ np.concatenate((H1r, G1r), axis=0)))
+
+ def forward(self, x):
+ B, N, c, k = x.shape # (B, N, k)
+ ns = math.floor(np.log2(N))
+ nl = pow(2, math.ceil(np.log2(N)))
+ extra_x = x[:, 0:nl - N, :, :]
+ x = torch.cat([x, extra_x], 1)
+ Ud = torch.jit.annotate(List[Tensor], [])
+ Us = torch.jit.annotate(List[Tensor], [])
+ # decompose
+ for i in range(ns - self.L):
+ # print('x shape',x.shape)
+ d, x = self.wavelet_transform(x)
+ Ud += [self.A(d) + self.B(x)]
+ Us += [self.C(d)]
+ x = self.T0(x) # coarsest scale transform
+
+ # reconstruct
+ for i in range(ns - 1 - self.L, -1, -1):
+ x = x + Us[i]
+ x = torch.cat((x, Ud[i]), -1)
+ x = self.evenOdd(x)
+ x = x[:, :N, :, :]
+
+ return x
+
+ def wavelet_transform(self, x):
+ xa = torch.cat([x[:, ::2, :, :],
+ x[:, 1::2, :, :],
+ ], -1)
+ d = torch.matmul(xa, self.ec_d)
+ s = torch.matmul(xa, self.ec_s)
+ return d, s
+
+ def evenOdd(self, x):
+
+ B, N, c, ich = x.shape # (B, N, c, k)
+ assert ich == 2 * self.k
+ x_e = torch.matmul(x, self.rc_e)
+ x_o = torch.matmul(x, self.rc_o)
+
+ x = torch.zeros(B, N * 2, c, self.k,
+ device=x.device)
+ x[..., ::2, :, :] = x_e
+ x[..., 1::2, :, :] = x_o
+ return x
diff --git a/basicts/archs/arch_zoo/fedformer_arch/self_attention_family.py b/basicts/archs/arch_zoo/fedformer_arch/self_attention_family.py
new file mode 100644
index 0000000..ea960bf
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/self_attention_family.py
@@ -0,0 +1,173 @@
+from math import sqrt
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+from .masking import TriangularCausalMask, ProbMask
+
+
+class FullAttention(nn.Module):
+ def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
+ super(FullAttention, self).__init__()
+ self.scale = scale
+ self.mask_flag = mask_flag
+ self.output_attention = output_attention
+ self.dropout = nn.Dropout(attention_dropout)
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, H, E = queries.shape
+ _, S, _, D = values.shape
+ scale = self.scale or 1. / sqrt(E)
+
+ scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+
+ if self.mask_flag:
+ if attn_mask is None:
+ attn_mask = TriangularCausalMask(B, L, device=queries.device)
+
+ scores.masked_fill_(attn_mask.mask, -np.inf)
+
+ A = self.dropout(torch.softmax(scale * scores, dim=-1))
+ V = torch.einsum("bhls,bshd->blhd", A, values)
+
+ if self.output_attention:
+ return (V.contiguous(), A)
+ else:
+ return (V.contiguous(), None)
+
+
+class ProbAttention(nn.Module):
+ def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
+ super(ProbAttention, self).__init__()
+ self.factor = factor
+ self.scale = scale
+ self.mask_flag = mask_flag
+ self.output_attention = output_attention
+ self.dropout = nn.Dropout(attention_dropout)
+
+ 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)
+ # real U = U_part(factor*ln(L_k))*L_q
+ index_sample = torch.randint(L_K, (L_Q, sample_k))
+ 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
+ if not self.mask_flag:
+ # V_sum = V.sum(dim=-2)
+ V_sum = V.mean(dim=-2)
+ contex = V_sum.unsqueeze(-2).expand(B, H,
+ L_Q, V_sum.shape[-1]).clone()
+ else: # use mask
+ # requires that L_Q == L_V, i.e. for self-attention only
+ assert (L_Q == L_V)
+ contex = V.cumsum(dim=-2)
+ return contex
+
+ def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
+ B, H, L_V, D = V.shape
+
+ if self.mask_flag:
+ attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
+ scores.masked_fill_(attn_mask.mask, -np.inf)
+
+ 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)
+ if self.output_attention:
+ attns = (torch.ones([B, H, L_V, L_V]) /
+ L_V).type_as(attn).to(attn.device)
+ attns[torch.arange(B)[:, None, None], torch.arange(H)[
+ None, :, None], index, :] = attn
+ return (context_in, attns)
+ else:
+ return (context_in, None)
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L_Q, H, D = queries.shape
+ _, L_K, _, _ = keys.shape
+
+ queries = queries.transpose(2, 1)
+ keys = keys.transpose(2, 1)
+ values = values.transpose(2, 1)
+
+ U_part = self.factor * \
+ np.ceil(np.log(L_K)).astype('int').item() # c*ln(L_k)
+ u = self.factor * \
+ np.ceil(np.log(L_Q)).astype('int').item() # c*ln(L_q)
+
+ U_part = U_part if U_part < L_K else L_K
+ u = u if u < L_Q else L_Q
+
+ scores_top, index = self._prob_QK(
+ queries, keys, sample_k=U_part, n_top=u)
+
+ # add scale factor
+ scale = self.scale or 1. / sqrt(D)
+ if scale is not None:
+ scores_top = scores_top * scale
+ # get the context
+ context = self._get_initial_context(values, L_Q)
+ # update the context with selected top_k queries
+ context, attn = self._update_context(
+ context, values, scores_top, index, L_Q, attn_mask)
+
+ return context.contiguous(), attn
+
+
+class AttentionLayer(nn.Module):
+ def __init__(self, attention, d_model, n_heads, d_keys=None,
+ d_values=None):
+ super(AttentionLayer, self).__init__()
+
+ d_keys = d_keys or (d_model // n_heads)
+ d_values = d_values or (d_model // n_heads)
+
+ self.inner_attention = attention
+ self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+ self.value_projection = nn.Linear(d_model, d_values * n_heads)
+ self.out_projection = nn.Linear(d_values * n_heads, d_model)
+ self.n_heads = n_heads
+
+ def forward(self, queries, keys, values, attn_mask):
+ B, L, _ = queries.shape
+ _, S, _ = keys.shape
+ H = self.n_heads
+
+ queries = self.query_projection(queries).view(B, L, H, -1)
+ keys = self.key_projection(keys).view(B, S, H, -1)
+ values = self.value_projection(values).view(B, S, H, -1)
+
+ out, attn = self.inner_attention(
+ queries,
+ keys,
+ values,
+ attn_mask
+ )
+ out = out.view(B, L, -1)
+
+ return self.out_projection(out), attn
diff --git a/basicts/archs/arch_zoo/fedformer_arch/utils.py b/basicts/archs/arch_zoo/fedformer_arch/utils.py
new file mode 100644
index 0000000..abad383
--- /dev/null
+++ b/basicts/archs/arch_zoo/fedformer_arch/utils.py
@@ -0,0 +1,393 @@
+import torch
+import torch.nn as nn
+
+import numpy as np
+from functools import partial
+
+from scipy.special import eval_legendre
+from sympy import Poly, legendre, Symbol, chebyshevt
+
+
+def legendreDer(k, x):
+ def _legendre(k, x):
+ return (2*k+1) * eval_legendre(k, x)
+ out = 0
+ for i in np.arange(k-1,-1,-2):
+ out += _legendre(i, x)
+ return out
+
+
+def phi_(phi_c, x, lb = 0, ub = 1):
+ mask = np.logical_or(x