[ { "path": "README.md", "content": "# MSGNet (AAAI2024)\n\nPaper Link:[MSGNet: Learning Multi-Scale Inter-Series Correlations for Multivariate Time Series Forecasting](https://arxiv.org/abs/2401.00423)\n\n## Usage\n\n- Train and evaluate MSGNet\n - You can use the following command:`sh ./scripts/ETTh1.sh`.\n\n- Train your model\n - Add model file in the folder `./models/your_model.py`.\n - Add model in the ***class*** Exp_Main.\n\n- Flight dataset\n - You can obtain the dataset from [Google Drive](https://drive.google.com/drive/folders/1JSZByfM0Ghat3g_D3a-puTZ2JsfebNWL?usp=sharing). Then please place it in the folder `./dataset`.\n\n## Model\n\nMSGNet employs several ScaleGraph blocks, each encompassing three pivotal modules: an FFT module for multi-scale data identification, an adaptive graph convolution module for inter-series correlation learning within a time scale, and a multi-head attention module for intra-series correlation learning.\n\n
\n \n
\n\n## Main Results\n\nForecast results with 96 review window and prediction length {96, 192, 336, 720}. The best result is represented in bold, followed by underline.\n\n
\n\n
\n\n## Citation\n\n```\n@article{cai2023msgnet,\n title={MSGNet: Learning Multi-Scale Inter-Series Correlations for Multivariate Time Series Forecasting},\n author={Cai, Wanlin and Liang, Yuxuan and Liu, Xianggen and Feng, Jianshuai and Wu, Yuankai},\n journal={arXiv preprint arXiv:2401.00423},\n year={2023}\n}\n```\n\n## Acknowledgement\n\nWe appreciate the valuable contributions of the following GitHub.\n\n- LTSF-Linear (https://github.com/cure-lab/LTSF-Linear)\n- TimesNet (https://github.com/thuml/TimesNet)\n- Time-Series-Library (https://github.com/thuml/Time-Series-Library)\n- MTGnn (https://github.com/nnzhan/MTGNN)\n- Autoformer (https://github.com/thuml/Autoformer)\n- Informer (https://github.com/zhouhaoyi/Informer2020)\n" }, { "path": "data_provider/data_factory.py", "content": "from .data_loader import Dataset_ETT_hour, Dataset_ETT_minute, \\\n Dataset_Custom, Dataset_Pred,Dataset_Flight\nfrom torch.utils.data import DataLoader\n\ndata_dict = {\n 'ETTh1': Dataset_ETT_hour,\n 'ETTh2': Dataset_ETT_hour,\n 'ETTm1': Dataset_ETT_minute,\n 'ETTm2': Dataset_ETT_minute,\n 'custom': Dataset_Custom,\n 'Flight':Dataset_Flight,\n}\n\n\n# flag = 'train' or 'val' or 'test'\ndef data_provider(args, flag):\n Data = data_dict[args.data]\n #time features encoding, options: [timeF, fixed, learned]\n timeenc = 0 if args.embed != 'timeF' else 1\n\n if flag == 'test':\n shuffle_flag = False\n drop_last = True\n batch_size = args.batch_size\n freq = args.freq\n\n elif flag == 'pred':\n shuffle_flag = False\n drop_last = False\n batch_size = 1\n freq = args.freq\n Data = Dataset_Pred\n\n else:\n shuffle_flag = True\n drop_last = True\n batch_size = args.batch_size\n freq = args.freq\n\n data_set = Data(\n root_path=args.root_path,\n data_path=args.data_path,\n flag=flag,\n size=[args.seq_len, args.label_len, args.pred_len],\n features=args.features,\n target=args.target,\n timeenc=timeenc,\n freq=freq,\n seasonal_patterns = args.seasonal_patterns\n )\n print(flag, len(data_set))\n data_loader = DataLoader(\n data_set,\n batch_size=batch_size,\n shuffle=shuffle_flag,\n num_workers=args.num_workers,\n drop_last=drop_last)\n return data_set, data_loader\n" }, { "path": "data_provider/data_loader.py", "content": "import os\nimport numpy as np\nimport pandas as pd\nimport os\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.preprocessing import StandardScaler\nfrom utils.timefeatures import time_features\nimport warnings\n\nwarnings.filterwarnings('ignore')\n\n#for Flight 4:4:2 split\nclass Dataset_Flight(Dataset):\n def __init__(self, root_path, flag='train', size=None,\n features='S', data_path='Flight.csv',\n target='OT', scale=True, timeenc=0, freq='h',seasonal_patterns=None):\n # size [seq_len, label_len, pred_len]\n # info\n if size == None:\n self.seq_len = 24 * 4 * 4\n self.label_len = 24 * 4\n self.pred_len = 24 * 4\n else:\n self.seq_len = size[0]\n self.label_len = size[1]\n self.pred_len = size[2]\n # init\n assert flag in ['train', 'test', 'val']\n type_map = {'train': 0, 'val': 1, 'test': 2}\n self.set_type = type_map[flag]\n\n self.features = features\n self.target = target\n self.scale = scale\n self.timeenc = timeenc\n self.freq = freq\n\n self.root_path = root_path\n self.data_path = data_path\n self.__read_data__()\n\n def __read_data__(self):\n self.scaler = StandardScaler()\n df_raw = pd.read_csv(os.path.join(self.root_path,\n self.data_path))\n\n '''\n df_raw.columns: ['date', ...(other features), target feature]\n '''\n cols = list(df_raw.columns)\n cols.remove(self.target)\n cols.remove('date')\n df_raw = df_raw[['date'] + cols + [self.target]]\n # print(cols)\n num_train = int(len(df_raw) * 0.4)\n num_test = int(len(df_raw) * 0.2)\n num_vali = len(df_raw) - num_train - num_test\n border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]\n border2s = [num_train, num_train + num_vali, len(df_raw)]\n border1 = border1s[self.set_type]\n border2 = border2s[self.set_type]\n\n if self.features == 'M' or self.features == 'MS':\n cols_data = df_raw.columns[1:]\n df_data = df_raw[cols_data]\n elif self.features == 'S':\n df_data = df_raw[[self.target]]\n\n if self.scale:\n train_data = df_data[border1s[0]:border2s[0]]\n self.scaler.fit(train_data.values)\n data = self.scaler.transform(df_data.values)\n else:\n data = df_data.values\n\n df_stamp = df_raw[['date']][border1:border2]\n df_stamp['date'] = pd.to_datetime(df_stamp.date)\n if self.timeenc == 0:\n df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)\n df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)\n df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)\n df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)\n data_stamp = df_stamp.drop(['date'], 1).values\n elif self.timeenc == 1:\n data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)\n data_stamp = data_stamp.transpose(1, 0)\n\n self.data_x = data[border1:border2]\n self.data_y = data[border1:border2]\n self.data_stamp = data_stamp\n\n def __getitem__(self, index):\n s_begin = index\n s_end = s_begin + self.seq_len\n r_begin = s_end - self.label_len\n r_end = r_begin + self.label_len + self.pred_len\n\n seq_x = self.data_x[s_begin:s_end]\n seq_y = self.data_y[r_begin:r_end]\n seq_x_mark = self.data_stamp[s_begin:s_end]\n seq_y_mark = self.data_stamp[r_begin:r_end]\n\n return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n def __len__(self):\n return len(self.data_x) - self.seq_len - self.pred_len + 1\n\n def inverse_transform(self, data):\n return self.scaler.inverse_transform(data)\n\nclass Dataset_Custom(Dataset):\n def __init__(self, root_path, flag='train', size=None,\n features='S', data_path='ETTh1.csv',\n target='OT', scale=True, timeenc=0, freq='h', seasonal_patterns=None):\n # size [seq_len, label_len, pred_len]\n # info\n if size == None:\n self.seq_len = 24 * 4 * 4\n self.label_len = 24 * 4\n self.pred_len = 24 * 4\n else:\n self.seq_len = size[0]\n self.label_len = size[1]\n self.pred_len = size[2]\n # init\n assert flag in ['train', 'test', 'val']\n type_map = {'train': 0, 'val': 1, 'test': 2}\n self.set_type = type_map[flag]\n\n self.features = features\n self.target = target\n self.scale = scale\n self.timeenc = timeenc\n self.freq = freq\n\n self.root_path = root_path\n self.data_path = data_path\n self.__read_data__()\n\n def __read_data__(self):\n self.scaler = StandardScaler()\n df_raw = pd.read_csv(os.path.join(self.root_path,\n self.data_path))\n\n '''\n df_raw.columns: ['date', ...(other features), target feature]\n '''\n cols = list(df_raw.columns)\n cols.remove(self.target)\n cols.remove('date')\n df_raw = df_raw[['date'] + cols + [self.target]]\n # print(cols)\n num_train = int(len(df_raw) * 0.7)\n num_test = int(len(df_raw) * 0.2)\n num_vali = len(df_raw) - num_train - num_test\n border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]\n border2s = [num_train, num_train + num_vali, len(df_raw)]\n border1 = border1s[self.set_type]\n border2 = border2s[self.set_type]\n\n if self.features == 'M' or self.features == 'MS':\n cols_data = df_raw.columns[1:]\n df_data = df_raw[cols_data]\n elif self.features == 'S':\n df_data = df_raw[[self.target]]\n\n if self.scale:\n train_data = df_data[border1s[0]:border2s[0]]\n self.scaler.fit(train_data.values)\n # print(self.scaler.mean_)\n # exit()\n data = self.scaler.transform(df_data.values)\n else:\n data = df_data.values\n\n df_stamp = df_raw[['date']][border1:border2]\n df_stamp['date'] = pd.to_datetime(df_stamp.date)\n if self.timeenc == 0:\n df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)\n df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)\n df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)\n df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)\n data_stamp = df_stamp.drop(['date'], 1).values\n elif self.timeenc == 1:\n data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)\n data_stamp = data_stamp.transpose(1, 0)\n\n self.data_x = data[border1:border2]\n self.data_y = data[border1:border2]\n self.data_stamp = data_stamp\n\n def __getitem__(self, index):\n s_begin = index\n s_end = s_begin + self.seq_len\n r_begin = s_end - self.label_len\n r_end = r_begin + self.label_len + self.pred_len\n\n seq_x = self.data_x[s_begin:s_end]\n seq_y = self.data_y[r_begin:r_end]\n seq_x_mark = self.data_stamp[s_begin:s_end]\n seq_y_mark = self.data_stamp[r_begin:r_end]\n\n return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n def __len__(self):\n return len(self.data_x) - self.seq_len - self.pred_len + 1\n\n def inverse_transform(self, data):\n return self.scaler.inverse_transform(data)\n\nclass Dataset_Pred(Dataset):\n def __init__(self, root_path, flag='pred', size=None,\n features='S', data_path='ETTh1.csv',\n target='OT', scale=True, inverse=False, timeenc=0, freq='15min', seasonal_patterns=None,cols=None):\n # size [seq_len, label_len, pred_len]\n # info\n if size == None:\n self.seq_len = 24 * 4 * 4\n self.label_len = 24 * 4\n self.pred_len = 24 * 4\n else:\n self.seq_len = size[0]\n self.label_len = size[1]\n self.pred_len = size[2]\n # init\n assert flag in ['pred']\n\n self.features = features\n self.target = target\n self.scale = scale\n self.inverse = inverse\n self.timeenc = timeenc\n self.freq = freq\n self.cols = cols\n self.root_path = root_path\n self.data_path = data_path\n self.__read_data__()\n\n def __read_data__(self):\n self.scaler = StandardScaler()\n df_raw = pd.read_csv(os.path.join(self.root_path,\n self.data_path))\n '''\n df_raw.columns: ['date', ...(other features), target feature]\n '''\n if self.cols:\n cols = self.cols.copy()\n cols.remove(self.target)\n else:\n cols = list(df_raw.columns)\n cols.remove(self.target)\n cols.remove('date')\n df_raw = df_raw[['date'] + cols + [self.target]]\n border1 = len(df_raw) - self.seq_len\n border2 = len(df_raw)\n\n if self.features == 'M' or self.features == 'MS':\n cols_data = df_raw.columns[1:]\n df_data = df_raw[cols_data]\n elif self.features == 'S':\n df_data = df_raw[[self.target]]\n if self.scale:\n self.scaler.fit(df_data.values)\n data = self.scaler.transform(df_data.values)\n else:\n data = df_data.values\n\n tmp_stamp = df_raw[['date']][border1:border2]\n tmp_stamp['date'] = pd.to_datetime(tmp_stamp.date)\n pred_dates = pd.date_range(tmp_stamp.date.values[-1], periods=self.pred_len + 1, freq=self.freq)\n\n df_stamp = pd.DataFrame(columns=['date'])\n df_stamp.date = list(tmp_stamp.date.values) + list(pred_dates[1:])\n if self.timeenc == 0:\n df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)\n df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)\n df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)\n df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)\n df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)\n df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)\n data_stamp = df_stamp.drop(['date'], 1).values\n elif self.timeenc == 1:\n data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)\n data_stamp = data_stamp.transpose(1, 0)\n\n self.data_x = data[border1:border2]\n if self.inverse:\n self.data_y = df_data.values[border1:border2]\n else:\n self.data_y = data[border1:border2]\n self.data_stamp = data_stamp\n\n def __getitem__(self, index):\n s_begin = index\n s_end = s_begin + self.seq_len\n r_begin = s_end - self.label_len\n r_end = r_begin + self.label_len + self.pred_len\n\n seq_x = self.data_x[s_begin:s_end]\n if self.inverse:\n seq_y = self.data_x[r_begin:r_begin + self.label_len]\n else:\n seq_y = self.data_y[r_begin:r_begin + self.label_len]\n seq_x_mark = self.data_stamp[s_begin:s_end]\n seq_y_mark = self.data_stamp[r_begin:r_end]\n\n return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n def __len__(self):\n return len(self.data_x) - self.seq_len + 1\n\n def inverse_transform(self, data):\n return self.scaler.inverse_transform(data)\n\nclass Dataset_ETT_hour(Dataset):\n def __init__(self, root_path, flag='train', size=None,\n features='S', data_path='ETTh1.csv',\n target='OT', scale=True, timeenc=0, freq='h', seasonal_patterns=None):\n # size [seq_len, label_len, pred_len]\n # info\n if size == None:\n self.seq_len = 24 * 4 * 4\n self.label_len = 24 * 4\n self.pred_len = 24 * 4\n else:\n self.seq_len = size[0]\n self.label_len = size[1]\n self.pred_len = size[2]\n # init\n assert flag in ['train', 'test', 'val']\n type_map = {'train': 0, 'val': 1, 'test': 2}\n self.set_type = type_map[flag]\n\n # M:multivariate predict multivariate, S:univariate predict univariate, MS:multivariate predict univariate'\n self.features = features\n self.target = target\n self.scale = scale\n self.timeenc = timeenc\n self.freq = freq\n\n self.root_path = root_path\n self.data_path = data_path\n self.__read_data__()\n\n def __read_data__(self):\n self.scaler = StandardScaler()\n df_raw = pd.read_csv(os.path.join(self.root_path,\n self.data_path))\n border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]\n border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]\n border1 = border1s[self.set_type]\n border2 = border2s[self.set_type]\n if self.features == 'M' or self.features == 'MS':\n cols_data = df_raw.columns[1:]\n df_data = df_raw[cols_data]\n elif self.features == 'S':\n df_data = df_raw[[self.target]]\n if self.scale:\n train_data = df_data[border1s[0]:border2s[0]]\n self.scaler.fit(train_data.values)\n data = self.scaler.transform(df_data.values)\n else:\n data = df_data.values\n df_stamp = df_raw[['date']][border1:border2]\n df_stamp['date'] = pd.to_datetime(df_stamp.date)\n if self.timeenc == 0:\n df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)\n df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)\n df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)\n df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)\n data_stamp = df_stamp.drop(['date'], 1).values\n elif self.timeenc == 1:\n data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)\n data_stamp = data_stamp.transpose(1, 0)\n\n self.data_x = data[border1:border2]\n self.data_y = data[border1:border2]\n self.data_stamp = data_stamp\n\n def __getitem__(self, index):\n s_begin = index\n s_end = s_begin + self.seq_len\n r_begin = s_end - self.label_len\n r_end = r_begin + self.label_len + self.pred_len\n\n seq_x = self.data_x[s_begin:s_end]\n seq_y = self.data_y[r_begin:r_end]\n seq_x_mark = self.data_stamp[s_begin:s_end]\n seq_y_mark = self.data_stamp[r_begin:r_end]\n\n return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n def __len__(self):\n return len(self.data_x) - self.seq_len - self.pred_len + 1\n\n def inverse_transform(self, data):\n return self.scaler.inverse_transform(data)\n\nclass Dataset_ETT_minute(Dataset):\n def __init__(self, root_path, flag='train', size=None,\n features='S', data_path='ETTm1.csv',\n target='OT', scale=True, timeenc=0, freq='t', seasonal_patterns=None):\n # size [seq_len, label_len, pred_len]\n # info\n if size == None:\n self.seq_len = 24 * 4 * 4\n self.label_len = 24 * 4\n self.pred_len = 24 * 4\n else:\n self.seq_len = size[0]\n self.label_len = size[1]\n self.pred_len = size[2]\n # init\n assert flag in ['train', 'test', 'val']\n type_map = {'train': 0, 'val': 1, 'test': 2}\n self.set_type = type_map[flag]\n\n self.features = features\n self.target = target\n self.scale = scale\n self.timeenc = timeenc\n self.freq = freq\n\n self.root_path = root_path\n self.data_path = data_path\n self.__read_data__()\n\n def __read_data__(self):\n self.scaler = StandardScaler()\n df_raw = pd.read_csv(os.path.join(self.root_path,\n self.data_path))\n\n border1s = [0,\n 12 * 30 * 24 * 4 - self.seq_len,\n 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]\n border2s = [12 * 30 * 24 * 4,\n 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4,\n 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]\n border1 = border1s[self.set_type]\n border2 = border2s[self.set_type]\n\n if self.features == 'M' or self.features == 'MS':\n cols_data = df_raw.columns[1:]\n df_data = df_raw[cols_data]\n elif self.features == 'S':\n df_data = df_raw[[self.target]]\n if self.scale:\n train_data = df_data[border1s[0]:border2s[0]]\n self.scaler.fit(train_data.values)\n data = self.scaler.transform(df_data.values)\n else:\n data = df_data.values\n\n df_stamp = df_raw[['date']][border1:border2]\n df_stamp['date'] = pd.to_datetime(df_stamp.date)\n if self.timeenc == 0:\n df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)\n df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)\n df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)\n df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)\n df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)\n df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)\n data_stamp = df_stamp.drop(['date'], 1).values\n elif self.timeenc == 1:\n data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)\n data_stamp = data_stamp.transpose(1, 0)\n\n self.data_x = data[border1:border2]\n self.data_y = data[border1:border2]\n self.data_stamp = data_stamp\n\n def __getitem__(self, index):\n s_begin = index\n s_end = s_begin + self.seq_len\n r_begin = s_end - self.label_len\n r_end = r_begin + self.label_len + self.pred_len\n\n seq_x = self.data_x[s_begin:s_end]\n seq_y = self.data_y[r_begin:r_end]\n seq_x_mark = self.data_stamp[s_begin:s_end]\n seq_y_mark = self.data_stamp[r_begin:r_end]\n\n return seq_x, seq_y, seq_x_mark, seq_y_mark\n\n def __len__(self):\n return len(self.data_x) - self.seq_len - self.pred_len + 1\n\n def inverse_transform(self, data):\n return self.scaler.inverse_transform(data)" }, { "path": "exp/exp_basic.py", "content": "import os\nimport torch\nimport numpy as np\n\n\nclass Exp_Basic(object):\n def __init__(self, args):\n self.args = args\n self.device = self._acquire_device()\n self.model = self._build_model().to(self.device)\n\n def _build_model(self):\n raise NotImplementedError\n return None\n\n def _acquire_device(self):\n if self.args.use_gpu:\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(\n self.args.gpu) if not self.args.use_multi_gpu else self.args.devices\n device = torch.device('cuda:{}'.format(self.args.gpu))\n print('Use GPU: cuda:{}'.format(self.args.gpu))\n else:\n device = torch.device('cpu')\n print('Use CPU')\n return device\n\n def _get_data(self):\n pass\n\n def vali(self):\n pass\n\n def train(self):\n pass\n\n def test(self):\n pass\n" }, { "path": "exp/exp_main.py", "content": "from data_provider.data_factory import data_provider\nfrom .exp_basic import Exp_Basic\nfrom models import Informer, Autoformer, DLinear, MSGNet\nfrom utils.tools import EarlyStopping, adjust_learning_rate, visual, test_params_flop\nfrom utils.metrics import metric\n\nimport torch\nimport torch.nn as nn\nfrom torch import optim, autograd\n\nimport os\nimport time\n\nimport warnings\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nwarnings.filterwarnings('ignore')\n\nclass Exp_Main(Exp_Basic):\n def __init__(self, args):\n super(Exp_Main, self).__init__(args)\n\n def _build_model(self):\n model_dict = {\n 'Informer': Informer,\n 'Autoformer': Autoformer,\n 'DLinear': DLinear,\n 'MSGNet': MSGNet\n }\n model = model_dict[self.args.model].Model(self.args).float()\n\n if self.args.use_multi_gpu and self.args.use_gpu:\n model = nn.DataParallel(model, device_ids=self.args.device_ids)\n return model\n\n #flag = 'train' or 'val' or 'test'\n def _get_data(self, flag):\n data_set, data_loader = data_provider(self.args, flag)\n return data_set, data_loader\n\n def _select_optimizer(self):\n model_optim = optim.Adam(self.model.parameters(), lr=self.args.learning_rate)\n return model_optim\n\n def _select_criterion(self):\n criterion = nn.MSELoss()\n return criterion\n\n\n def vali(self, vali_data, vali_loader, criterion):\n total_loss = []\n self.model.eval()\n with torch.no_grad():\n for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(vali_loader):\n batch_x = batch_x.float().to(self.device)\n batch_y = batch_y.float()\n batch_x_mark = batch_x_mark.float().to(self.device)\n batch_y_mark = batch_y_mark.float().to(self.device)\n\n # decoder input\n dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()\n dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)\n # encoder - decoder\n if self.args.use_amp:\n with torch.cuda.amp.autocast():\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n else:\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n f_dim = -1 if self.args.features == 'MS' else 0\n outputs = outputs[:, -self.args.pred_len:, f_dim:]\n batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)\n\n pred = outputs.detach().cpu()\n true = batch_y.detach().cpu()\n\n loss = criterion(pred, true)\n\n total_loss.append(loss)\n total_loss = np.average(total_loss)\n self.model.train()\n return total_loss\n\n def train(self, setting):\n train_data, train_loader = self._get_data(flag='train')\n vali_data, vali_loader = self._get_data(flag='val')\n test_data, test_loader = self._get_data(flag='test')\n\n path = os.path.join(self.args.checkpoints, setting)\n if not os.path.exists(path):\n os.makedirs(path)\n\n time_now = time.time()\n train_steps = len(train_loader)\n early_stopping = EarlyStopping(patience=self.args.patience, verbose=True)\n\n model_optim = self._select_optimizer()\n criterion = self._select_criterion()\n #use automatic mixed precision training\n if self.args.use_amp:\n scaler = torch.cuda.amp.GradScaler()\n for epoch in range(self.args.train_epochs):\n iter_count = 0\n train_loss = []\n\n self.model.train()\n epoch_time = time.time()\n for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(train_loader):\n\n iter_count += 1\n model_optim.zero_grad()\n\n batch_x = batch_x.float().to(self.device)\n batch_y = batch_y.float().to(self.device)\n batch_x_mark = batch_x_mark.float().to(self.device)\n batch_y_mark = batch_y_mark.float().to(self.device)\n\n dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()\n dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)\n\n # encoder - decoder\n if self.args.use_amp:\n with torch.cuda.amp.autocast():\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n\n f_dim = -1 if self.args.features == 'MS' else 0\n outputs = outputs[:, -self.args.pred_len:, f_dim:]\n batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)\n loss = criterion(outputs, batch_y)\n train_loss.append(loss.item())\n else:\n if 'Linear' in self.args.model:\n # print(\"Linear\")\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention: #whether to output attention in ecoder\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n \n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n # print(outputs.shape,batch_y.shape)\n f_dim = -1 if self.args.features == 'MS' else 0\n outputs = outputs[:, -self.args.pred_len:, f_dim:]\n batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)\n loss = criterion(outputs, batch_y)\n train_loss.append(loss.item())\n\n if (i + 1) % 100 == 0:\n print(\"\\titers: {0}, epoch: {1} | loss: {2:.7f}\".format(i + 1, epoch + 1, loss.item()))\n speed = (time.time() - time_now) / iter_count\n left_time = speed * ((self.args.train_epochs - epoch) * train_steps - i)\n print('\\tspeed: {:.4f}s/iter; left time: {:.4f}s'.format(speed, left_time))\n iter_count = 0\n time_now = time.time()\n\n if self.args.use_amp:\n scaler.scale(loss).backward()\n scaler.step(model_optim)\n scaler.update()\n else:\n with autograd.detect_anomaly():\n loss.backward()\n model_optim.step()\n\n print(\"Epoch: {} cost time: {}\".format(epoch + 1, time.time() - epoch_time))\n train_loss = np.average(train_loss)\n vali_loss = self.vali(vali_data, vali_loader, criterion)\n test_loss = self.vali(test_data, test_loader, criterion)\n\n print(\"Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}\".format(\n epoch + 1, train_steps, train_loss, vali_loss, test_loss))\n early_stopping(vali_loss, self.model, path)\n if early_stopping.early_stop:\n print(\"Early stopping\")\n break\n\n adjust_learning_rate(model_optim, epoch + 1, self.args)\n\n best_model_path = path + '/' + 'checkpoint.pth'\n self.model.load_state_dict(torch.load(best_model_path))\n\n return self.model\n\n def test(self, setting, test=0):\n test_data, test_loader = self._get_data(flag='test')\n if test:\n print('loading model')\n self.model.load_state_dict(torch.load(os.path.join('./checkpoints/' + setting, 'checkpoint.pth')))\n\n preds = []\n trues = []\n inputx = []\n folder_path = './test_results/' + setting + '/'\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n\n self.model.eval()\n with torch.no_grad():\n for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(test_loader):\n\n batch_x = batch_x.float().to(self.device)\n batch_y = batch_y.float().to(self.device)\n batch_x_mark = batch_x_mark.float().to(self.device)\n batch_y_mark = batch_y_mark.float().to(self.device)\n\n # decoder input\n dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()\n dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)\n # encoder - decoder\n if self.args.use_amp:\n with torch.cuda.amp.autocast():\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n else:\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n\n f_dim = -1 if self.args.features == 'MS' else 0\n # print(outputs.shape,batch_y.shape)\n outputs = outputs[:, -self.args.pred_len:, f_dim:]\n batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)\n outputs = outputs.detach().cpu().numpy()\n batch_y = batch_y.detach().cpu().numpy()\n\n pred = outputs # outputs.detach().cpu().numpy() # .squeeze()\n true = batch_y # batch_y.detach().cpu().numpy() # .squeeze()\n\n preds.append(pred)\n trues.append(true)\n inputx.append(batch_x.detach().cpu().numpy())\n if i % 10 == 0:\n input = batch_x.detach().cpu().numpy()\n gt = np.concatenate((input[0, :, -1], true[0, :, -1]), axis=0)\n pd = np.concatenate((input[0, :, -1], pred[0, :, -1]), axis=0)\n visual(gt, pd, os.path.join(folder_path, str(i) + '.pdf'))\n #See utils / tools for usage\n if self.args.test_flop:\n test_params_flop((batch_x.shape[1],batch_x.shape[2]))\n exit()\n # print('preds_shape:', len(preds),len(preds[0]),len(preds[1]))\n\n preds = np.array(preds)\n trues = np.array(trues)\n inputx = np.array(inputx)\n\n print('preds_shape:', preds.shape)\n print('trues_shape:', trues.shape)\n\n preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])\n trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])\n inputx = inputx.reshape(-1, inputx.shape[-2], inputx.shape[-1])\n\n # result save\n folder_path = './results/' + setting + '/'\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n\n mae, mse, rmse, mape, mspe, rse, corr, nd, nrmse = metric(preds, trues)\n print('nd:{}, nrmse:{}, mse:{}, mae:{}, rse:{}, mape:{}'.format(nd, nrmse,mse, mae, rse, mape))\n f = open(\"result.txt\", 'a')\n f.write(setting + \" \\n\")\n f.write('nd:{}, nrmse:{}, mse:{}, mae:{}, rse:{}, mape:{}'.format(nd, nrmse,mse, mae, rse, mape))\n f.write('\\n')\n f.write('\\n')\n f.close()\n\n # np.save(folder_path + 'metrics.npy', np.array([mae, mse, rmse, mape, mspe,rse, corr]))\n # np.save(folder_path + 'pred.npy', preds)\n # np.save(folder_path + 'true.npy', trues)\n # np.save(folder_path + 'x.npy', inputx)\n return\n\n\n def predict(self, setting, load=False):\n pred_data, pred_loader = self._get_data(flag='pred')\n\n if load:\n path = os.path.join(self.args.checkpoints, setting)\n best_model_path = path + '/' + 'checkpoint.pth'\n self.model.load_state_dict(torch.load(best_model_path))\n\n preds = []\n\n self.model.eval()\n with torch.no_grad():\n for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(pred_loader):\n batch_x = batch_x.float().to(self.device)\n batch_y = batch_y.float()\n batch_x_mark = batch_x_mark.float().to(self.device)\n batch_y_mark = batch_y_mark.float().to(self.device)\n\n # decoder input\n dec_inp = torch.zeros([batch_y.shape[0], self.args.pred_len, batch_y.shape[2]]).float().to(batch_y.device)\n dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)\n # encoder - decoder\n if self.args.use_amp:\n with torch.cuda.amp.autocast():\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n else:\n if 'Linear' in self.args.model:\n outputs = self.model(batch_x)\n else:\n if self.args.output_attention:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]\n else:\n outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)\n pred = outputs.detach().cpu().numpy() # .squeeze()\n preds.append(pred)\n\n preds = np.array(preds)\n preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])\n\n # result save\n folder_path = './results/' + setting + '/'\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n\n np.save(folder_path + 'real_prediction.npy', preds)\n\n return\n" }, { "path": "exp/exp_stat.py", "content": "from data_provider.data_factory import data_provider\nfrom exp.exp_basic import Exp_Basic\nfrom utils.tools import EarlyStopping, adjust_learning_rate, visual\nfrom utils.metrics import metric\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch import optim\n\nimport os\nimport time\nimport warnings\nimport matplotlib.pyplot as plt\nfrom models.Stat_models import *\n\nwarnings.filterwarnings('ignore')\n\n\nclass Exp_Main(Exp_Basic):\n def __init__(self, args):\n super(Exp_Main, self).__init__(args)\n\n def _build_model(self):\n model_dict = {\n 'Naive': Naive_repeat,\n 'ARIMA': Arima,\n 'SARIMA': SArima,\n 'GBRT': GBRT,\n }\n model = model_dict[self.args.model](self.args).float()\n\n return model\n\n def _get_data(self, flag):\n data_set, data_loader = data_provider(self.args, flag)\n return data_set, data_loader\n\n def test(self, setting, test=0):\n test_data, test_loader = self._get_data(flag='test')\n\n # Sample 10% \n samples = max(int(self.args.sample * self.args.batch_size),1)\n\n preds = []\n trues = []\n inputx = []\n folder_path = './test_results/' + setting + '/'\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n\n with torch.no_grad():\n for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(test_loader):\n batch_x = batch_x.float().to(self.device).cpu().numpy()\n batch_y = batch_y.float().to(self.device).cpu().numpy()\n\n batch_x = batch_x[:samples]\n outputs = self.model(batch_x)\n\n f_dim = -1 if self.args.features == 'MS' else 0\n # print(outputs.shape,batch_y.shape)\n outputs = outputs[:, -self.args.pred_len:, f_dim:]\n batch_y = batch_y[:samples, -self.args.pred_len:, f_dim:]\n\n pred = outputs # outputs.detach().cpu().numpy() # .squeeze()\n true = batch_y # batch_y.detach().cpu().numpy() # .squeeze()\n\n preds.append(pred)\n trues.append(true)\n inputx.append(batch_x)\n if i % 20 == 0:\n input = batch_x\n gt = np.concatenate((input[0, :, -1], true[0, :, -1]), axis=0)\n pd = np.concatenate((input[0, :, -1], pred[0, :, -1]), axis=0)\n visual(gt, pd, os.path.join(folder_path, str(i) + '.pdf'))\n\n preds = np.array(preds)\n trues = np.array(trues)\n inputx = np.array(inputx)\n\n preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])\n trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])\n inputx = inputx.reshape(-1, inputx.shape[-2], inputx.shape[-1])\n\n folder_path = './results/' + setting + '/'\n if not os.path.exists(folder_path):\n os.makedirs(folder_path)\n\n mae, mse, rmse, mape, mspe, rse, corr = metric(preds, trues)\n corr = []\n print('mse:{}, mae:{}, rse:{}, corr:{}'.format(mse, mae, rse, corr))\n f = open(\"result.txt\", 'a')\n f.write(setting + \" \\n\")\n f.write('mse:{}, mae:{}, rse:{}, corr:{}'.format(mse, mae, rse, corr))\n f.write('\\n')\n f.write('\\n')\n f.close()\n\n np.save(folder_path + 'metrics.npy', np.array([mae, mse, rmse, mape, mspe,rse, corr]))\n np.save(folder_path + 'pred.npy', preds)\n np.save(folder_path + 'true.npy', trues)\n # np.save(folder_path + 'x.npy', inputx)\n return" }, { "path": "layers/AutoCorrelation.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport math\nfrom math import sqrt\nimport os\n\n\nclass AutoCorrelation(nn.Module):\n \"\"\"\n AutoCorrelation Mechanism with the following two phases:\n (1) period-based dependencies discovery\n (2) time delay aggregation\n This block can replace the self-attention family mechanism seamlessly.\n \"\"\"\n def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False):\n super(AutoCorrelation, self).__init__()\n self.factor = factor\n self.scale = scale\n self.mask_flag = mask_flag\n self.output_attention = output_attention\n self.dropout = nn.Dropout(attention_dropout)\n\n def time_delay_agg_training(self, values, corr):\n \"\"\"\n SpeedUp version of Autocorrelation (a batch-normalization style design)\n This is for the training phase.\n \"\"\"\n head = values.shape[1]\n channel = values.shape[2]\n length = values.shape[3]\n # find top k\n top_k = int(self.factor * math.log(length))\n mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)\n index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]\n weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)\n # update corr\n tmp_corr = torch.softmax(weights, dim=-1)\n # aggregation\n tmp_values = values\n delays_agg = torch.zeros_like(values).float()\n for i in range(top_k):\n pattern = torch.roll(tmp_values, -int(index[i]), -1)\n delays_agg = delays_agg + pattern * \\\n (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))\n return delays_agg\n\n def time_delay_agg_inference(self, values, corr):\n \"\"\"\n SpeedUp version of Autocorrelation (a batch-normalization style design)\n This is for the inference phase.\n \"\"\"\n batch = values.shape[0]\n head = values.shape[1]\n channel = values.shape[2]\n length = values.shape[3]\n # index init\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n # init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()\n init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).to(device)\n\n # find top k\n top_k = int(self.factor * math.log(length))\n mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)\n weights = torch.topk(mean_value, top_k, dim=-1)[0]\n delay = torch.topk(mean_value, top_k, dim=-1)[1]\n # update corr\n tmp_corr = torch.softmax(weights, dim=-1)\n # aggregation\n tmp_values = values.repeat(1, 1, 1, 2)\n delays_agg = torch.zeros_like(values).float()\n for i in range(top_k):\n tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)\n pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)\n delays_agg = delays_agg + pattern * \\\n (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))\n return delays_agg\n\n def time_delay_agg_full(self, values, corr):\n \"\"\"\n Standard version of Autocorrelation\n \"\"\"\n batch = values.shape[0]\n head = values.shape[1]\n channel = values.shape[2]\n length = values.shape[3]\n # index init\n init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()\n # find top k\n top_k = int(self.factor * math.log(length))\n weights = torch.topk(corr, top_k, dim=-1)[0]\n delay = torch.topk(corr, top_k, dim=-1)[1]\n # update corr\n tmp_corr = torch.softmax(weights, dim=-1)\n # aggregation\n tmp_values = values.repeat(1, 1, 1, 2)\n delays_agg = torch.zeros_like(values).float()\n for i in range(top_k):\n tmp_delay = init_index + delay[..., i].unsqueeze(-1)\n pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)\n delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))\n return delays_agg\n\n def forward(self, queries, keys, values, attn_mask):\n B, L, H, E = queries.shape\n _, S, _, D = values.shape\n if L > S:\n zeros = torch.zeros_like(queries[:, :(L - S), :]).float()\n values = torch.cat([values, zeros], dim=1)\n keys = torch.cat([keys, zeros], dim=1)\n else:\n values = values[:, :L, :, :]\n keys = keys[:, :L, :, :]\n\n # period-based dependencies (b, len//period , period , d_model) ->(b, period ,d_model, len//period)\n #(b , T, h , n) ->(b, h, n, T)\n q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1)\n k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)\n #(b, period ,d_model, period),\n res = q_fft * torch.conj(k_fft)\n corr = torch.fft.irfft(res, dim=-1)\n\n # time delay agg\n if self.training:\n V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)\n else:\n V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)\n\n if self.output_attention:\n return (V.contiguous(), corr.permute(0, 3, 1, 2))\n else:\n return (V.contiguous(), None)\n\n\nclass AutoCorrelationLayer(nn.Module):\n def __init__(self, correlation, d_model, n_heads, d_keys=None,\n d_values=None):\n super(AutoCorrelationLayer, self).__init__()\n\n d_keys = d_keys or (d_model // n_heads)\n d_values = d_values or (d_model // n_heads)\n\n self.inner_correlation = correlation\n self.query_projection = nn.Linear(d_model, d_keys * n_heads)\n self.key_projection = nn.Linear(d_model, d_keys * n_heads)\n self.value_projection = nn.Linear(d_model, d_values * n_heads)\n self.out_projection = nn.Linear(d_values * n_heads, d_model)\n self.n_heads = n_heads\n\n def forward(self, queries, keys, values, attn_mask):\n B, L, _ = queries.shape\n _, S, _ = keys.shape\n H = self.n_heads\n\n queries = self.query_projection(queries).view(B, L, H, -1)\n keys = self.key_projection(keys).view(B, S, H, -1)\n values = self.value_projection(values).view(B, S, H, -1)\n\n out, attn = self.inner_correlation(\n queries,\n keys,\n values,\n attn_mask\n )\n out = out.view(B, L, -1)\n\n return self.out_projection(out), attn\n" }, { "path": "layers/Autoformer_EncDec.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass my_Layernorm(nn.Module):\n \"\"\"\n Special designed layernorm for the seasonal part\n \"\"\"\n def __init__(self, channels):\n super(my_Layernorm, self).__init__()\n self.layernorm = nn.LayerNorm(channels)\n\n def forward(self, x):\n x_hat = self.layernorm(x)\n bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)\n return x_hat - bias\n\n\nclass moving_avg(nn.Module):\n \"\"\"\n Moving average block to highlight the trend of time series\n \"\"\"\n def __init__(self, kernel_size, stride):\n super(moving_avg, self).__init__()\n self.kernel_size = kernel_size\n self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)\n\n def forward(self, x):\n # padding on the both ends of time series\n front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)\n end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)\n x = torch.cat([front, x, end], dim=1)\n x = self.avg(x.permute(0, 2, 1))\n x = x.permute(0, 2, 1)\n return x\n\n\nclass series_decomp(nn.Module):\n \"\"\"\n Series decomposition block\n \"\"\"\n def __init__(self, kernel_size):\n super(series_decomp, self).__init__()\n self.moving_avg = moving_avg(kernel_size, stride=1)\n\n def forward(self, x):\n moving_mean = self.moving_avg(x)\n res = x - moving_mean\n return res, moving_mean\n\n\nclass EncoderLayer(nn.Module):\n \"\"\"\n Autoformer encoder layer with the progressive decomposition architecture\n \"\"\"\n def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation=\"relu\"):\n super(EncoderLayer, self).__init__()\n d_ff = d_ff or 4 * d_model\n self.attention = attention\n self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)\n self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)\n self.decomp1 = series_decomp(moving_avg)\n self.decomp2 = series_decomp(moving_avg)\n self.dropout = nn.Dropout(dropout)\n self.activation = F.relu if activation == \"relu\" else F.gelu\n\n def forward(self, x, attn_mask=None):\n new_x, attn = self.attention(\n x, x, x,\n attn_mask=attn_mask\n )\n x = x + self.dropout(new_x)\n x, _ = self.decomp1(x)\n y = x\n y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n y = self.dropout(self.conv2(y).transpose(-1, 1))\n res, _ = self.decomp2(x + y)\n return res, attn\n\n\nclass Encoder(nn.Module):\n \"\"\"\n Autoformer encoder\n \"\"\"\n def __init__(self, attn_layers, conv_layers=None, norm_layer=None):\n super(Encoder, self).__init__()\n self.attn_layers = nn.ModuleList(attn_layers)\n self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None\n self.norm = norm_layer\n\n def forward(self, x, attn_mask=None):\n attns = []\n if self.conv_layers is not None:\n for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):\n x, attn = attn_layer(x, attn_mask=attn_mask)\n x = conv_layer(x)\n attns.append(attn)\n x, attn = self.attn_layers[-1](x)\n attns.append(attn)\n else:\n for attn_layer in self.attn_layers:\n x, attn = attn_layer(x, attn_mask=attn_mask)\n attns.append(attn)\n\n if self.norm is not None:\n x = self.norm(x)\n\n return x, attns\n\n\nclass DecoderLayer(nn.Module):\n \"\"\"\n Autoformer decoder layer with the progressive decomposition architecture\n \"\"\"\n def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,\n moving_avg=25, dropout=0.1, activation=\"relu\"):\n super(DecoderLayer, self).__init__()\n d_ff = d_ff or 4 * d_model\n self.self_attention = self_attention\n self.cross_attention = cross_attention\n self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)\n self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)\n self.decomp1 = series_decomp(moving_avg)\n self.decomp2 = series_decomp(moving_avg)\n self.decomp3 = series_decomp(moving_avg)\n self.dropout = nn.Dropout(dropout)\n self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,\n padding_mode='circular', bias=False)\n self.activation = F.relu if activation == \"relu\" else F.gelu\n\n def forward(self, x, cross, x_mask=None, cross_mask=None):\n x = x + self.dropout(self.self_attention(\n x, x, x,\n attn_mask=x_mask\n )[0])\n x, trend1 = self.decomp1(x)\n x = x + self.dropout(self.cross_attention(\n x, cross, cross,\n attn_mask=cross_mask\n )[0])\n x, trend2 = self.decomp2(x)\n y = x\n y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n y = self.dropout(self.conv2(y).transpose(-1, 1))\n x, trend3 = self.decomp3(x + y)\n\n residual_trend = trend1 + trend2 + trend3\n residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)\n return x, residual_trend\n\n\nclass Decoder(nn.Module):\n \"\"\"\n Autoformer encoder\n \"\"\"\n def __init__(self, layers, norm_layer=None, projection=None):\n super(Decoder, self).__init__()\n self.layers = nn.ModuleList(layers)\n self.norm = norm_layer\n self.projection = projection\n\n def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):\n for layer in self.layers:\n x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)\n trend = trend + residual_trend\n\n if self.norm is not None:\n x = self.norm(x)\n\n if self.projection is not None:\n x = self.projection(x)\n return x, trend\n" }, { "path": "layers/Embed.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn.utils import weight_norm\nimport math\n\n\nclass PositionalEmbedding(nn.Module):\n def __init__(self, d_model, max_len=5000):\n super(PositionalEmbedding, self).__init__()\n # Compute the positional encodings once in log space.\n pe = torch.zeros(max_len, d_model).float()\n pe.require_grad = False\n\n position = torch.arange(0, max_len).float().unsqueeze(1)\n div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n pe[:, 0::2] = torch.sin(position * div_term)\n pe[:, 1::2] = torch.cos(position * div_term)\n\n pe = pe.unsqueeze(0)\n self.register_buffer('pe', pe)\n\n def forward(self, x):\n return self.pe[:, :x.size(1)]\n\n\nclass TokenEmbedding(nn.Module):\n def __init__(self, c_in, d_model):\n super(TokenEmbedding, self).__init__()\n padding = 1 if torch.__version__ >= '1.5.0' else 2\n self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,\n kernel_size=3, padding=padding, padding_mode='circular', bias=False)\n for m in self.modules():\n if isinstance(m, nn.Conv1d):\n nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')\n\n def forward(self, x):\n x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)\n return x\n\n\nclass FixedEmbedding(nn.Module):\n def __init__(self, c_in, d_model):\n super(FixedEmbedding, self).__init__()\n\n w = torch.zeros(c_in, d_model).float()\n w.require_grad = False\n\n position = torch.arange(0, c_in).float().unsqueeze(1)\n div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()\n\n w[:, 0::2] = torch.sin(position * div_term)\n w[:, 1::2] = torch.cos(position * div_term)\n\n self.emb = nn.Embedding(c_in, d_model)\n self.emb.weight = nn.Parameter(w, requires_grad=False)\n\n def forward(self, x):\n return self.emb(x).detach()\n\n\nclass TemporalEmbedding(nn.Module):\n def __init__(self, d_model, embed_type='fixed', freq='h'):\n super(TemporalEmbedding, self).__init__()\n\n minute_size = 4\n hour_size = 24\n weekday_size = 7\n day_size = 32\n month_size = 13\n\n Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding\n if freq == 't':\n self.minute_embed = Embed(minute_size, d_model)\n self.hour_embed = Embed(hour_size, d_model)\n self.weekday_embed = Embed(weekday_size, d_model)\n self.day_embed = Embed(day_size, d_model)\n self.month_embed = Embed(month_size, d_model)\n\n def forward(self, x):\n x = x.long()\n\n minute_x = self.minute_embed(x[:, :, 4]) if hasattr(self, 'minute_embed') else 0.\n hour_x = self.hour_embed(x[:, :, 3])\n weekday_x = self.weekday_embed(x[:, :, 2])\n day_x = self.day_embed(x[:, :, 1])\n month_x = self.month_embed(x[:, :, 0])\n\n return hour_x + weekday_x + day_x + month_x + minute_x\n\nclass TimeFeatureEmbedding(nn.Module):\n def __init__(self, d_model, embed_type='timeF', freq='h'):\n super(TimeFeatureEmbedding, self).__init__()\n freq_map = {'h': 4, 't': 5, 's': 6, 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}\n d_inp = freq_map[freq]\n self.embed = nn.Linear(d_inp, d_model, bias=False)\n\n def forward(self, x):\n return self.embed(x)\n\n\nclass DataEmbedding(nn.Module):\n def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n super(DataEmbedding, self).__init__()\n self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n #(batch_size, len batch_x[1], d_model )\n self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,\n freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(\n d_model=d_model, embed_type=embed_type, freq=freq)\n #(1, len batch_x[1], d_model)\n self.position_embedding = PositionalEmbedding(d_model=d_model)\n\n self.dropout = nn.Dropout(p=dropout)\n\n def forward(self, x, x_mark):\n if x_mark is None:\n x = self.value_embedding(x) + self.position_embedding(x)\n else:\n x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)\n return self.dropout(x)\n\n\nclass DataEmbedding_wo_pos(nn.Module):\n def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n super(DataEmbedding_wo_pos, self).__init__()\n\n self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n self.position_embedding = PositionalEmbedding(d_model=d_model)\n self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,\n freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(\n d_model=d_model, embed_type=embed_type, freq=freq)\n self.dropout = nn.Dropout(p=dropout)\n\n def forward(self, x, x_mark):\n x = self.value_embedding(x) + self.temporal_embedding(x_mark)\n return self.dropout(x)\n\nclass DataEmbedding_wo_pos_temp(nn.Module):\n def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n super(DataEmbedding_wo_pos_temp, self).__init__()\n\n self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n self.position_embedding = PositionalEmbedding(d_model=d_model)\n self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,\n freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(\n d_model=d_model, embed_type=embed_type, freq=freq)\n self.dropout = nn.Dropout(p=dropout)\n\n def forward(self, x, x_mark):\n x = self.value_embedding(x)\n return self.dropout(x)\n\nclass DataEmbedding_wo_temp(nn.Module):\n def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):\n super(DataEmbedding_wo_temp, self).__init__()\n\n self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n self.position_embedding = PositionalEmbedding(d_model=d_model)\n self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,\n freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(\n d_model=d_model, embed_type=embed_type, freq=freq)\n self.dropout = nn.Dropout(p=dropout)\n\n def forward(self, x, x_mark):\n x = self.value_embedding(x) + self.position_embedding(x)\n return self.dropout(x)" }, { "path": "layers/MSGBlock.py", "content": "from math import sqrt\r\nimport numpy as np\r\nimport torch.nn as nn\r\nimport torch.nn.functional as F\r\nimport torch\r\nfrom torch import nn, Tensor\r\nfrom einops import rearrange\r\nfrom einops.layers.torch import Rearrange\r\nfrom utils.masking import TriangularCausalMask\r\n\r\nclass Predict(nn.Module):\r\n def __init__(self, individual, c_out, seq_len, pred_len, dropout):\r\n super(Predict, self).__init__()\r\n self.individual = individual\r\n self.c_out = c_out\r\n\r\n if self.individual:\r\n self.seq2pred = nn.ModuleList()\r\n self.dropout = nn.ModuleList()\r\n for i in range(self.c_out):\r\n self.seq2pred.append(nn.Linear(seq_len , pred_len))\r\n self.dropout.append(nn.Dropout(dropout))\r\n else:\r\n self.seq2pred = nn.Linear(seq_len , pred_len)\r\n self.dropout = nn.Dropout(dropout)\r\n\r\n #(B, c_out , seq)\r\n def forward(self, x):\r\n if self.individual:\r\n out = []\r\n for i in range(self.c_out):\r\n per_out = self.seq2pred[i](x[:,i,:])\r\n per_out = self.dropout[i](per_out)\r\n out.append(per_out)\r\n out = torch.stack(out,dim=1)\r\n else:\r\n out = self.seq2pred(x)\r\n out = self.dropout(out)\r\n\r\n return out\r\n\r\n\r\nclass Attention_Block(nn.Module):\r\n def __init__(self, d_model, d_ff=None, n_heads=8, dropout=0.1, activation=\"relu\"):\r\n super(Attention_Block, self).__init__()\r\n d_ff = d_ff or 4 * d_model\r\n self.attention = self_attention(FullAttention, d_model, n_heads=n_heads)\r\n self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\r\n self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\r\n self.norm1 = nn.LayerNorm(d_model)\r\n self.norm2 = nn.LayerNorm(d_model)\r\n self.dropout = nn.Dropout(dropout)\r\n self.activation = F.relu if activation == \"relu\" else F.gelu\r\n\r\n def forward(self, x, attn_mask=None):\r\n new_x, attn = self.attention(\r\n x, x, x,\r\n attn_mask=attn_mask\r\n )\r\n x = x + self.dropout(new_x)\r\n\r\n y = x = self.norm1(x)\r\n y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\r\n y = self.dropout(self.conv2(y).transpose(-1, 1))\r\n\r\n return self.norm2(x + y)\r\n\r\n\r\nclass self_attention(nn.Module):\r\n def __init__(self, attention, d_model ,n_heads):\r\n super(self_attention, self).__init__()\r\n d_keys = d_model // n_heads\r\n d_values = d_model // n_heads\r\n\r\n self.inner_attention = attention( attention_dropout = 0.1)\r\n self.query_projection = nn.Linear(d_model, d_keys * n_heads)\r\n self.key_projection = nn.Linear(d_model, d_keys * n_heads)\r\n self.value_projection = nn.Linear(d_model, d_values * n_heads)\r\n self.out_projection = nn.Linear(d_values * n_heads, d_model)\r\n self.n_heads = n_heads\r\n\r\n\r\n def forward(self, queries ,keys ,values, attn_mask= None):\r\n B, L, _ = queries.shape\r\n _, S, _ = keys.shape\r\n H = self.n_heads\r\n queries = self.query_projection(queries).view(B, L, H, -1)\r\n keys = self.key_projection(keys).view(B, S, H, -1)\r\n values = self.value_projection(values).view(B, S, H, -1)\r\n\r\n out, attn = self.inner_attention(\r\n queries,\r\n keys,\r\n values,\r\n attn_mask\r\n )\r\n out = out.view(B, L, -1)\r\n out = self.out_projection(out)\r\n return out , attn\r\n\r\n\r\nclass FullAttention(nn.Module):\r\n def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\r\n super(FullAttention, self).__init__()\r\n self.scale = scale\r\n self.mask_flag = mask_flag\r\n self.output_attention = output_attention\r\n self.dropout = nn.Dropout(attention_dropout)\r\n\r\n def forward(self, queries, keys, values, attn_mask):\r\n B, L, H, E = queries.shape\r\n _, S, _, D = values.shape\r\n scale = self.scale or 1. / sqrt(E)\r\n scores = torch.einsum(\"blhe,bshe->bhls\", queries, keys)\r\n if self.mask_flag:\r\n if attn_mask is None:\r\n attn_mask = TriangularCausalMask(B, L, device=queries.device)\r\n scores.masked_fill_(attn_mask.mask, -np.inf)\r\n A = self.dropout(torch.softmax(scale * scores, dim=-1))\r\n V = torch.einsum(\"bhls,bshd->blhd\", A, values)\r\n # return V.contiguous()\r\n if self.output_attention:\r\n return (V.contiguous(), A)\r\n else:\r\n return (V.contiguous(), None)\r\n\r\n\r\nclass GraphBlock(nn.Module):\r\n def __init__(self, c_out , d_model , conv_channel, skip_channel,\r\n gcn_depth , dropout, propalpha ,seq_len , node_dim):\r\n super(GraphBlock, self).__init__()\r\n\r\n self.nodevec1 = nn.Parameter(torch.randn(c_out, node_dim), requires_grad=True)\r\n self.nodevec2 = nn.Parameter(torch.randn(node_dim, c_out), requires_grad=True)\r\n self.start_conv = nn.Conv2d(1 , conv_channel, (d_model - c_out + 1, 1))\r\n self.gconv1 = mixprop(conv_channel, skip_channel, gcn_depth, dropout, propalpha)\r\n self.gelu = nn.GELU()\r\n self.end_conv = nn.Conv2d(skip_channel, seq_len , (1, seq_len ))\r\n self.linear = nn.Linear(c_out, d_model)\r\n self.norm = nn.LayerNorm(d_model)\r\n\r\n # x in (B, T, d_model)\r\n # Here we use a mlp to fit a complex mapping f (x)\r\n def forward(self, x):\r\n adp = F.softmax(F.relu(torch.mm(self.nodevec1, self.nodevec2)), dim=1)\r\n out = x.unsqueeze(1).transpose(2, 3)\r\n out = self.start_conv(out)\r\n out = self.gelu(self.gconv1(out , adp))\r\n out = self.end_conv(out).squeeze()\r\n out = self.linear(out)\r\n\r\n return self.norm(x + out)\r\n\r\n\r\nclass nconv(nn.Module):\r\n def __init__(self):\r\n super(nconv,self).__init__()\r\n\r\n def forward(self,x, A):\r\n x = torch.einsum('ncwl,vw->ncvl',(x,A))\r\n # x = torch.einsum('ncwl,wv->nclv',(x,A)\r\n return x.contiguous()\r\n\r\n\r\nclass linear(nn.Module):\r\n def __init__(self,c_in,c_out,bias=True):\r\n super(linear,self).__init__()\r\n self.mlp = torch.nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0,0), stride=(1,1), bias=bias)\r\n\r\n def forward(self,x):\r\n return self.mlp(x)\r\n\r\n\r\nclass mixprop(nn.Module):\r\n def __init__(self,c_in,c_out,gdep,dropout,alpha):\r\n super(mixprop, self).__init__()\r\n self.nconv = nconv()\r\n self.mlp = linear((gdep+1)*c_in,c_out)\r\n self.gdep = gdep\r\n self.dropout = dropout\r\n self.alpha = alpha\r\n\r\n def forward(self, x, adj):\r\n adj = adj + torch.eye(adj.size(0)).to(x.device)\r\n d = adj.sum(1)\r\n h = x\r\n out = [h]\r\n a = adj / d.view(-1, 1)\r\n for i in range(self.gdep):\r\n h = self.alpha*x + (1-self.alpha)*self.nconv(h,a)\r\n out.append(h)\r\n ho = torch.cat(out,dim=1)\r\n ho = self.mlp(ho)\r\n return ho\r\n\r\n\r\nclass simpleVIT(nn.Module):\r\n def __init__(self, in_channels, emb_size, patch_size=2, depth=1, num_heads=4, dropout=0.1,init_weight =True):\r\n super(simpleVIT, self).__init__()\r\n self.emb_size = emb_size\r\n self.depth = depth\r\n self.to_patch = nn.Sequential(\r\n nn.Conv2d(in_channels, emb_size, 2 * patch_size + 1, padding= patch_size),\r\n Rearrange('b e (h) (w) -> b (h w) e'),\r\n )\r\n self.layers = nn.ModuleList([])\r\n for _ in range(self.depth):\r\n self.layers.append(nn.ModuleList([\r\n nn.LayerNorm(emb_size),\r\n MultiHeadAttention(emb_size, num_heads, dropout),\r\n FeedForward(emb_size, emb_size)\r\n ]))\r\n\r\n if init_weight:\r\n self._initialize_weights()\r\n\r\n def _initialize_weights(self):\r\n for m in self.modules():\r\n if isinstance(m, nn.Conv2d):\r\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\r\n if m.bias is not None:\r\n nn.init.constant_(m.bias, 0)\r\n\r\n def forward(self,x):\r\n B , N ,_ ,P = x.shape\r\n x = self.to_patch(x)\r\n # x = x.permute(0, 2, 3, 1).reshape(B,-1, N)\r\n for norm ,attn, ff in self.layers:\r\n x = attn(norm(x)) + x\r\n x = ff(x) + x\r\n\r\n x = x.transpose(1,2).reshape(B, self.emb_size ,-1, P)\r\n return x\r\n\r\nclass MultiHeadAttention(nn.Module):\r\n def __init__(self, emb_size, num_heads, dropout):\r\n super().__init__()\r\n self.emb_size = emb_size\r\n self.num_heads = num_heads\r\n self.keys = nn.Linear(emb_size, emb_size)\r\n self.queries = nn.Linear(emb_size, emb_size)\r\n self.values = nn.Linear(emb_size, emb_size)\r\n self.att_drop = nn.Dropout(dropout)\r\n self.projection = nn.Linear(emb_size, emb_size)\r\n\r\n def forward(self, x: Tensor, mask: Tensor = None) -> Tensor:\r\n queries = rearrange(self.queries(x), \"b n (h d) -> b h n d\", h=self.num_heads)\r\n keys = rearrange(self.keys(x), \"b n (h d) -> b h n d\", h=self.num_heads)\r\n values = rearrange(self.values(x), \"b n (h d) -> b h n d\", h=self.num_heads)\r\n energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)\r\n if mask is not None:\r\n fill_value = torch.finfo(torch.float32).min\r\n energy.mask_fill(~mask, fill_value)\r\n\r\n scaling = self.emb_size ** (1 / 2)\r\n att = F.softmax(energy, dim=-1) / scaling\r\n att = self.att_drop(att)\r\n # sum up over the third axis\r\n out = torch.einsum('bhal, bhlv -> bhav ', att, values)\r\n out = rearrange(out, \"b h n d -> b n (h d)\")\r\n out = self.projection(out)\r\n return out\r\n\r\nclass FeedForward(nn.Module):\r\n def __init__(self, dim, hidden_dim):\r\n super().__init__()\r\n self.net = nn.Sequential(\r\n nn.LayerNorm(dim),\r\n nn.Linear(dim, hidden_dim),\r\n nn.GELU(),\r\n nn.Linear(hidden_dim, dim),\r\n )\r\n def forward(self, x):\r\n return self.net(x)" }, { "path": "layers/SelfAttention_Family.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport matplotlib.pyplot as plt\n\nimport numpy as np\nimport math\nfrom math import sqrt\nfrom utils.masking import TriangularCausalMask, ProbMask\nimport os\n\nclass FullAttention(nn.Module):\n def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n super(FullAttention, self).__init__()\n self.scale = scale\n self.mask_flag = mask_flag\n self.output_attention = output_attention\n self.dropout = nn.Dropout(attention_dropout)\n\n def forward(self, queries, keys, values, attn_mask):\n B, L, H, E = queries.shape\n _, S, _, D = values.shape\n scale = self.scale or 1. / sqrt(E)\n scores = torch.einsum(\"blhe,bshe->bhls\", queries, keys)\n if self.mask_flag:\n if attn_mask is None:\n attn_mask = TriangularCausalMask(B, L, device=queries.device)\n scores.masked_fill_(attn_mask.mask, -np.inf)\n A = self.dropout(torch.softmax(scale * scores, dim=-1))\n V = torch.einsum(\"bhls,bshd->blhd\", A, values)\n if self.output_attention:\n return (V.contiguous(), A)\n else:\n return (V.contiguous(), None)\n\n\nclass ProbAttention(nn.Module):\n def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):\n super(ProbAttention, self).__init__()\n self.factor = factor\n self.scale = scale\n self.mask_flag = mask_flag\n self.output_attention = output_attention\n self.dropout = nn.Dropout(attention_dropout)\n\n def _prob_QK(self, Q, K, sample_k, n_top): # n_top: c*ln(L_q)\n # Q [B, H, L, D]\n B, H, L_K, E = K.shape\n _, _, L_Q, _ = Q.shape\n\n # calculate the sampled Q_K\n K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)\n index_sample = torch.randint(L_K, (L_Q, sample_k)) # real U = U_part(factor*ln(L_k))*L_q\n K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]\n Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()\n\n # find the Top_k query with sparisty measurement\n M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)\n M_top = M.topk(n_top, sorted=False)[1]\n\n # use the reduced Q to calculate Q_K\n Q_reduce = Q[torch.arange(B)[:, None, None],\n torch.arange(H)[None, :, None],\n M_top, :] # factor*ln(L_q)\n Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) # factor*ln(L_q)*L_k\n return Q_K, M_top\n\n def _get_initial_context(self, V, L_Q):\n B, H, L_V, D = V.shape\n if not self.mask_flag:\n # V_sum = V.sum(dim=-2)\n V_sum = V.mean(dim=-2)\n contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()\n else: # use mask\n assert (L_Q == L_V) # requires that L_Q == L_V, i.e. for self-attention only\n contex = V.cumsum(dim=-2)\n return contex\n\n def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):\n B, H, L_V, D = V.shape\n\n if self.mask_flag:\n attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)\n scores.masked_fill_(attn_mask.mask, -np.inf)\n\n attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores)\n\n context_in[torch.arange(B)[:, None, None],\n torch.arange(H)[None, :, None],\n index, :] = torch.matmul(attn, V).type_as(context_in)\n if self.output_attention:\n attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)\n attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn\n return (context_in, attns)\n else:\n return (context_in, None)\n\n def forward(self, queries, keys, values, attn_mask):\n B, L_Q, H, D = queries.shape\n _, L_K, _, _ = keys.shape\n\n queries = queries.transpose(2, 1)\n keys = keys.transpose(2, 1)\n values = values.transpose(2, 1)\n\n U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item() # c*ln(L_k)\n u = self.factor * np.ceil(np.log(L_Q)).astype('int').item() # c*ln(L_q)\n\n U_part = U_part if U_part < L_K else L_K\n u = u if u < L_Q else L_Q\n\n scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)\n\n # add scale factor\n scale = self.scale or 1. / sqrt(D)\n if scale is not None:\n scores_top = scores_top * scale\n # get the context\n context = self._get_initial_context(values, L_Q)\n # update the context with selected top_k queries\n context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)\n\n return context.contiguous(), attn\n\n\nclass AttentionLayer(nn.Module):\n def __init__(self, attention, d_model, n_heads, d_keys=None,\n d_values=None):\n super(AttentionLayer, self).__init__()\n\n d_keys = d_keys or (d_model // n_heads)\n d_values = d_values or (d_model // n_heads)\n\n self.inner_attention = attention\n self.query_projection = nn.Linear(d_model, d_keys * n_heads)\n self.key_projection = nn.Linear(d_model, d_keys * n_heads)\n self.value_projection = nn.Linear(d_model, d_values * n_heads)\n self.out_projection = nn.Linear(d_values * n_heads, d_model)\n self.n_heads = n_heads\n\n def forward(self, queries, keys, values, attn_mask):\n B, L, _ = queries.shape\n _, S, _ = keys.shape\n H = self.n_heads\n queries = self.query_projection(queries).view(B, L, H, -1)\n keys = self.key_projection(keys).view(B, S, H, -1)\n values = self.value_projection(values).view(B, S, H, -1)\n out, attn = self.inner_attention(\n queries,\n keys,\n values,\n attn_mask\n )\n out = out.view(B, L, -1)\n\n return self.out_projection(out), attn\n" }, { "path": "layers/Transformer_EncDec.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass ConvLayer(nn.Module):\n def __init__(self, c_in):\n super(ConvLayer, self).__init__()\n self.downConv = nn.Conv1d(in_channels=c_in,\n out_channels=c_in,\n kernel_size=3,\n padding=2,\n padding_mode='circular')\n self.norm = nn.BatchNorm1d(c_in)\n self.activation = nn.ELU()\n self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)\n\n def forward(self, x):\n x = self.downConv(x.permute(0, 2, 1))\n x = self.norm(x)\n x = self.activation(x)\n x = self.maxPool(x)\n x = x.transpose(1, 2)\n return x\n\nclass EncoderLayer(nn.Module):\n def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation=\"relu\"):\n super(EncoderLayer, self).__init__()\n d_ff = d_ff or 4 * d_model\n self.attention = attention\n self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n self.norm1 = nn.LayerNorm(d_model)\n self.norm2 = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(dropout)\n self.activation = F.relu if activation == \"relu\" else F.gelu\n def forward(self, x, attn_mask=None):\n new_x, attn = self.attention(\n x, x, x,\n attn_mask=attn_mask\n )\n x = x + self.dropout(new_x)\n\n y = x = self.norm1(x)\n y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n y = self.dropout(self.conv2(y).transpose(-1, 1))\n\n return self.norm2(x + y), attn\n\n\nclass Encoder(nn.Module):\n def __init__(self, attn_layers, conv_layers=None, norm_layer=None):\n super(Encoder, self).__init__()\n self.attn_layers = nn.ModuleList(attn_layers)\n self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None\n self.norm = norm_layer\n\n def forward(self, x, attn_mask=None):\n attns = []\n if self.conv_layers is not None:\n for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):\n x, attn = attn_layer(x, attn_mask=attn_mask)\n x = conv_layer(x)\n attns.append(attn)\n x, attn = self.attn_layers[-1](x)\n attns.append(attn)\n else:\n for attn_layer in self.attn_layers:\n x, attn = attn_layer(x, attn_mask=attn_mask)\n attns.append(attn)\n\n if self.norm is not None:\n x = self.norm(x)\n\n return x, attns\n\n\nclass DecoderLayer(nn.Module):\n def __init__(self, self_attention, cross_attention, d_model, d_ff=None,\n dropout=0.1, activation=\"relu\"):\n super(DecoderLayer, self).__init__()\n d_ff = d_ff or 4 * d_model\n self.self_attention = self_attention\n self.cross_attention = cross_attention\n self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n self.norm1 = nn.LayerNorm(d_model)\n self.norm2 = nn.LayerNorm(d_model)\n self.norm3 = nn.LayerNorm(d_model)\n self.dropout = nn.Dropout(dropout)\n self.activation = F.relu if activation == \"relu\" else F.gelu\n # dec_out(in), enc_out\n def forward(self, x, cross, x_mask=None, cross_mask=None):\n x = x + self.dropout(self.self_attention(\n x, x, x,\n attn_mask=x_mask\n )[0])\n x = self.norm1(x)\n\n x = x + self.dropout(self.cross_attention(\n x, cross, cross, #q,k,v\n attn_mask=cross_mask\n )[0])\n\n y = x = self.norm2(x)\n y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))\n y = self.dropout(self.conv2(y).transpose(-1, 1))\n\n return self.norm3(x + y)\n\n\nclass Decoder(nn.Module):\n def __init__(self, layers, norm_layer=None, projection=None):\n super(Decoder, self).__init__()\n self.layers = nn.ModuleList(layers)\n self.norm = norm_layer\n self.projection = projection\n\n #self.decoder(dec_out(in), enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)\n def forward(self, x, cross, x_mask=None, cross_mask=None ,external=None):\n for layer in self.layers:\n x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)\n\n if self.norm is not None:\n x = self.norm(x)\n\n if self.projection is not None:\n x = self.projection(x)\n return x\n" }, { "path": "models/Autoformer.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom layers.Embed import DataEmbedding, DataEmbedding_wo_pos,DataEmbedding_wo_pos_temp,DataEmbedding_wo_temp\nfrom layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer\nfrom layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp\nimport math\nimport numpy as np\n\n\nclass Model(nn.Module):\n \"\"\"\n Autoformer is the first method to achieve the series-wise connection,\n with inherent O(LlogL) complexity\n \"\"\"\n def __init__(self, configs):\n super(Model, self).__init__()\n self.seq_len = configs.seq_len\n self.label_len = configs.label_len\n self.pred_len = configs.pred_len\n self.output_attention = configs.output_attention\n\n # Decomp\n kernel_size = configs.moving_avg\n self.decomp = series_decomp(kernel_size)\n\n # Embedding\n # The series-wise connection inherently contains the sequential information.\n # Thus, we can discard the position embedding of transformers.\n if configs.embed_type == 0:\n self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 1:\n self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 2:\n self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n\n elif configs.embed_type == 3:\n self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 4:\n self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n \n # Encoder\n self.encoder = Encoder(\n [\n EncoderLayer(\n AutoCorrelationLayer(\n AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,\n output_attention=configs.output_attention),\n configs.d_model, configs.n_heads),\n configs.d_model,\n configs.d_ff,\n moving_avg=configs.moving_avg,\n dropout=configs.dropout,\n activation=configs.activation\n ) for l in range(configs.e_layers)\n ],\n norm_layer=my_Layernorm(configs.d_model)\n )\n # Decoder\n self.decoder = Decoder(\n [\n DecoderLayer(\n AutoCorrelationLayer(\n AutoCorrelation(True, configs.factor, attention_dropout=configs.dropout,\n output_attention=False),\n configs.d_model, configs.n_heads),\n AutoCorrelationLayer(\n AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,\n output_attention=False),\n configs.d_model, configs.n_heads),\n configs.d_model,\n configs.c_out,\n configs.d_ff,\n moving_avg=configs.moving_avg,\n dropout=configs.dropout,\n activation=configs.activation,\n )\n for l in range(configs.d_layers)\n ],\n norm_layer=my_Layernorm(configs.d_model),\n projection=nn.Linear(configs.d_model, configs.c_out, bias=True)\n )\n\n def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,\n enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):\n # decomp init\n mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)\n zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)\n seasonal_init, trend_init = self.decomp(x_enc)\n # decoder input\n trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)\n seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)\n # enc\n enc_out = self.enc_embedding(x_enc, x_mark_enc)\n enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)\n # dec\n dec_out = self.dec_embedding(seasonal_init, x_mark_dec)\n seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,\n trend=trend_init)\n # final\n dec_out = trend_part + seasonal_part\n\n if self.output_attention:\n return dec_out[:, -self.pred_len:, :], attns\n else:\n return dec_out[:, -self.pred_len:, :] # [B, L, D]\n" }, { "path": "models/DLinear.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\n\nclass moving_avg(nn.Module):\n \"\"\"\n Moving average block to highlight the trend of time series\n \"\"\"\n def __init__(self, kernel_size, stride):\n super(moving_avg, self).__init__()\n self.kernel_size = kernel_size\n self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)\n\n def forward(self, x):\n # padding on the both ends of time series\n front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)\n end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)\n x = torch.cat([front, x, end], dim=1)\n # print(x.permute(0, 2, 1))\n x = self.avg(x.permute(0, 2, 1))\n x = x.permute(0, 2, 1)\n return x\n\n\nclass series_decomp(nn.Module):\n \"\"\"\n Series decomposition block\n \"\"\"\n def __init__(self, kernel_size):\n super(series_decomp, self).__init__()\n self.moving_avg = moving_avg(kernel_size, stride=1)\n\n def forward(self, x):\n moving_mean = self.moving_avg(x)\n res = x - moving_mean\n return res, moving_mean\n\nclass Model(nn.Module):\n \"\"\"\n Decomposition-Linear\n \"\"\"\n def __init__(self, configs):\n super(Model, self).__init__()\n self.seq_len = configs.seq_len\n self.pred_len = configs.pred_len\n\n # Decompsition Kernel Size\n kernel_size = 25\n self.decompsition = series_decomp(kernel_size) #return res, moving_mean\n self.individual = configs.individual\n self.channels = configs.enc_in\n\n if self.individual:\n self.Linear_Seasonal = nn.ModuleList()\n self.Linear_Trend = nn.ModuleList()\n \n for i in range(self.channels):\n self.Linear_Seasonal.append(nn.Linear(self.seq_len,self.pred_len))\n self.Linear_Trend.append(nn.Linear(self.seq_len,self.pred_len))\n\n # Use this two lines if you want to visualize the weights\n # self.Linear_Seasonal[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))\n # self.Linear_Trend[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))\n else:\n self.Linear_Seasonal = nn.Linear(self.seq_len,self.pred_len)\n self.Linear_Trend = nn.Linear(self.seq_len,self.pred_len)\n \n # Use this two lines if you want to visualize the weights\n # self.Linear_Seasonal.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))\n # self.Linear_Trend.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))\n\n def forward(self, x):\n # x: [Batch, Input length, Channel]\n seasonal_init, trend_init = self.decompsition(x) #return res, moving_mean\n seasonal_init, trend_init = seasonal_init.permute(0,2,1), trend_init.permute(0,2,1)\n if self.individual:\n seasonal_output = torch.zeros([seasonal_init.size(0),seasonal_init.size(1),self.pred_len],dtype=seasonal_init.dtype).to(seasonal_init.device)\n trend_output = torch.zeros([trend_init.size(0),trend_init.size(1),self.pred_len],dtype=trend_init.dtype).to(trend_init.device)\n for i in range(self.channels):\n seasonal_output[:,i,:] = self.Linear_Seasonal[i](seasonal_init[:,i,:])\n trend_output[:,i,:] = self.Linear_Trend[i](trend_init[:,i,:])\n else:\n seasonal_output = self.Linear_Seasonal(seasonal_init)\n trend_output = self.Linear_Trend(trend_init)\n\n x = seasonal_output + trend_output\n return x.permute(0,2,1) # to [Batch, Output length, Channel]\n" }, { "path": "models/Informer.py", "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom utils.masking import TriangularCausalMask, ProbMask\nfrom layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer\nfrom layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer\nfrom layers.Embed import DataEmbedding,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp\nimport numpy as np\n\n\nclass Model(nn.Module):\n \"\"\"\n Informer with Propspare attention in O(LlogL) complexity\n \"\"\"\n def __init__(self, configs):\n super(Model, self).__init__()\n self.pred_len = configs.pred_len\n self.output_attention = configs.output_attention\n\n # Embedding\n if configs.embed_type == 0:\n self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 1:\n self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 2:\n self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n\n elif configs.embed_type == 3:\n self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n elif configs.embed_type == 4:\n self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,\n configs.dropout)\n # Encoder\n self.encoder = Encoder(\n [\n EncoderLayer(\n AttentionLayer(\n ProbAttention(False, configs.factor, attention_dropout=configs.dropout,\n output_attention=configs.output_attention),\n configs.d_model, configs.n_heads),\n configs.d_model,\n configs.d_ff,\n dropout=configs.dropout,\n activation=configs.activation\n ) for l in range(configs.e_layers)\n ],\n [\n ConvLayer(\n configs.d_model\n ) for l in range(configs.e_layers - 1)\n ] if configs.distil else None,\n norm_layer=torch.nn.LayerNorm(configs.d_model)\n )\n # Decoder\n self.decoder = Decoder(\n [\n DecoderLayer(\n AttentionLayer(\n ProbAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),\n configs.d_model, configs.n_heads),\n AttentionLayer(\n ProbAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),\n configs.d_model, configs.n_heads),\n configs.d_model,\n configs.d_ff,\n dropout=configs.dropout,\n activation=configs.activation,\n )\n for l in range(configs.d_layers)\n ],\n norm_layer=torch.nn.LayerNorm(configs.d_model),\n projection=nn.Linear(configs.d_model, configs.c_out, bias=True)\n )\n\n def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,\n enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):\n\n enc_out = self.enc_embedding(x_enc, x_mark_enc)\n enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)\n\n dec_out = self.dec_embedding(x_dec, x_mark_dec)\n dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)\n\n if self.output_attention:\n return dec_out[:, -self.pred_len:, :], attns\n else:\n return dec_out[:, -self.pred_len:, :] # [B, L, D]\n" }, { "path": "models/MSGNet.py", "content": "import numpy as np\r\n# import pywt\r\nimport torch\r\nimport torch.nn as nn\r\nimport torch.nn.functional as F\r\nimport torch.fft\r\nfrom layers.Embed import DataEmbedding\r\nfrom layers.MSGBlock import GraphBlock, simpleVIT, Attention_Block, Predict\r\n\r\n\r\ndef FFT_for_Period(x, k=2):\r\n # [B, T, C]\r\n xf = torch.fft.rfft(x, dim=1)\r\n frequency_list = abs(xf).mean(0).mean(-1)\r\n frequency_list[0] = 0\r\n _, top_list = torch.topk(frequency_list, k)\r\n top_list = top_list.detach().cpu().numpy()\r\n period = x.shape[1] // top_list\r\n return period, abs(xf).mean(-1)[:, top_list]\r\n\r\n\r\nclass ScaleGraphBlock(nn.Module):\r\n def __init__(self, configs):\r\n super(ScaleGraphBlock, self).__init__()\r\n self.seq_len = configs.seq_len\r\n self.pred_len = configs.pred_len\r\n self.k = configs.top_k\r\n\r\n self.att0 = Attention_Block(configs.d_model, configs.d_ff,\r\n n_heads=configs.n_heads, dropout=configs.dropout, activation=\"gelu\")\r\n self.norm = nn.LayerNorm(configs.d_model)\r\n self.gelu = nn.GELU()\r\n self.gconv = nn.ModuleList()\r\n for i in range(self.k):\r\n self.gconv.append(\r\n GraphBlock(configs.c_out , configs.d_model , configs.conv_channel, configs.skip_channel,\r\n configs.gcn_depth , configs.dropout, configs.propalpha ,configs.seq_len,\r\n configs.node_dim))\r\n\r\n\r\n def forward(self, x):\r\n B, T, N = x.size()\r\n scale_list, scale_weight = FFT_for_Period(x, self.k)\r\n res = []\r\n for i in range(self.k):\r\n scale = scale_list[i]\r\n #Gconv\r\n x = self.gconv[i](x)\r\n # paddng\r\n if (self.seq_len) % scale != 0:\r\n length = (((self.seq_len) // scale) + 1) * scale\r\n padding = torch.zeros([x.shape[0], (length - (self.seq_len)), x.shape[2]]).to(x.device)\r\n out = torch.cat([x, padding], dim=1)\r\n else:\r\n length = self.seq_len\r\n out = x\r\n out = out.reshape(B, length // scale, scale, N)\r\n\r\n #for Mul-attetion\r\n out = out.reshape(-1 , scale , N)\r\n out = self.norm(self.att0(out))\r\n out = self.gelu(out)\r\n out = out.reshape(B, -1 , scale , N).reshape(B ,-1 ,N)\r\n # #for simpleVIT\r\n # out = self.att(out.permute(0, 3, 1, 2).contiguous()) #return\r\n # out = out.permute(0, 2, 3, 1).reshape(B, -1 ,N)\r\n\r\n out = out[:, :self.seq_len, :]\r\n res.append(out)\r\n\r\n res = torch.stack(res, dim=-1)\r\n # adaptive aggregation\r\n scale_weight = F.softmax(scale_weight, dim=1)\r\n scale_weight = scale_weight.unsqueeze(1).unsqueeze(1).repeat(1, T, N, 1)\r\n res = torch.sum(res * scale_weight, -1)\r\n # residual connection\r\n res = res + x\r\n return res\r\n\r\n\r\nclass Model(nn.Module):\r\n def __init__(self, configs):\r\n super(Model, self).__init__()\r\n self.configs = configs\r\n self.task_name = configs.task_name\r\n self.seq_len = configs.seq_len\r\n self.label_len = configs.label_len\r\n self.pred_len = configs.pred_len\r\n self.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\r\n\r\n # for graph\r\n # self.num_nodes = configs.c_out\r\n # self.subgraph_size = configs.subgraph_size\r\n # self.node_dim = configs.node_dim\r\n # to return adj (node , node)\r\n # self.graph = constructor_graph()\r\n\r\n self.model = nn.ModuleList([ScaleGraphBlock(configs) for _ in range(configs.e_layers)])\r\n self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model,\r\n configs.embed, configs.freq, configs.dropout)\r\n self.layer = configs.e_layers\r\n self.layer_norm = nn.LayerNorm(configs.d_model)\r\n self.predict_linear = nn.Linear(\r\n self.seq_len, self.pred_len + self.seq_len)\r\n self.projection = nn.Linear(\r\n configs.d_model, configs.c_out, bias=True)\r\n self.seq2pred = Predict(configs.individual ,configs.c_out,\r\n configs.seq_len, configs.pred_len, configs.dropout)\r\n\r\n\r\n def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):\r\n # Normalization from Non-stationary Transformer\r\n means = x_enc.mean(1, keepdim=True).detach()\r\n x_enc = x_enc - means\r\n stdev = torch.sqrt(\r\n torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5)\r\n x_enc /= stdev\r\n\r\n # embedding\r\n enc_out = self.enc_embedding(x_enc, x_mark_enc) # [B,T,C]\r\n # adp = self.graph(torch.arange(self.num_nodes).to(self.device))\r\n for i in range(self.layer):\r\n enc_out = self.layer_norm(self.model[i](enc_out))\r\n\r\n # porject back\r\n dec_out = self.projection(enc_out)\r\n dec_out = self.seq2pred(dec_out.transpose(1, 2)).transpose(1, 2)\r\n\r\n # De-Normalization from Non-stationary Transformer\r\n dec_out = dec_out * \\\r\n (stdev[:, 0, :].unsqueeze(1).repeat(\r\n 1, self.pred_len, 1))\r\n dec_out = dec_out + \\\r\n (means[:, 0, :].unsqueeze(1).repeat(\r\n 1, self.pred_len, 1))\r\n\r\n return dec_out[:, -self.pred_len:, :]\r\n\r\n\r\n" }, { "path": "run_longExp.py", "content": "import argparse\nimport os\nimport time\nfrom multiprocessing import freeze_support\nimport torch\nfrom exp.exp_main import Exp_Main\nimport random\nimport numpy as np\n\nfix_seed = 2021\nrandom.seed(fix_seed)\ntorch.manual_seed(fix_seed)\nnp.random.seed(fix_seed)\n\nparser = argparse.ArgumentParser(description='MSGNet for Time Series Forecasting')\n\n# basic config\nparser.add_argument('--task_name', type=str, required=False, default='long_term_forecast',\n help='task name, options:[long_term_forecast, mask, short_term_forecast, imputation, classification, anomaly_detection]')\nparser.add_argument('--is_training', type=int, required=True, default=1, help='status')\nparser.add_argument('--model_id', type=str, required=True, default='test', help='model id')\nparser.add_argument('--model', type=str, required=True, default='Autoformer',\n help='model name, options: [Autoformer, Informer, Transformer]')\n\n# data loader\nparser.add_argument('--data', type=str, required=True, default='ETTm1', help='dataset type')\nparser.add_argument('--root_path', type=str, default='./data/ETT/', help='root path of the data file')\nparser.add_argument('--data_path', type=str, default='ETTh1.csv', help='data file')\nparser.add_argument('--features', type=str, default='M',\n help='forecasting task, options:[M, S, MS]; M:multivariate predict multivariate,'\n ' S:univariate predict univariate, MS:multivariate predict univariate')\nparser.add_argument('--target', type=str, default='OT', help='target feature in S or MS task')\nparser.add_argument('--freq', type=str, default='h',\n help='freq for time features encoding, '\n 'options:[s:secondly, t:minutely, h:hourly, d:daily, b:business days, w:weekly, m:monthly], '\n 'you can also use more detailed freq like 15min or 3h')\nparser.add_argument('--checkpoints', type=str, default='./checkpoints/', help='location of model checkpoints')\n\n# forecasting task\nparser.add_argument('--seq_len', type=int, default=96, help='input sequence length')\nparser.add_argument('--label_len', type=int, default=48, help='start token length')\nparser.add_argument('--pred_len', type=int, default=96, help='prediction sequence length')\nparser.add_argument('--seasonal_patterns', type=str, default='Monthly', help='subset for M4')\n\n\nparser.add_argument('--top_k', type=int, default=5, help='for TimesBlock/ScaleGraphBlock')\nparser.add_argument('--num_kernels', type=int, default=6, help='for Inception')\n\nparser.add_argument('--num_nodes', type=int, default=7, help='to create Graph')\nparser.add_argument('--subgraph_size', type=int, default=3, help='neighbors number')\nparser.add_argument('--tanhalpha', type=float, default=3, help='')\n\n#GCN\nparser.add_argument('--node_dim', type=int, default=10, help='each node embbed to dim dimentions')\nparser.add_argument('--gcn_depth', type=int, default=2, help='')\nparser.add_argument('--gcn_dropout', type=float, default=0.3, help='')\nparser.add_argument('--propalpha', type=float, default=0.3, help='')\nparser.add_argument('--conv_channel', type=int, default=32, help='')\nparser.add_argument('--skip_channel', type=int, default=32, help='')\n\n\n# DLinear\nparser.add_argument('--individual', action='store_true', default=False, help='DLinear: a linear layer for each variate(channel) individually')\n# Formers\nparser.add_argument('--embed_type', type=int, default=0, help='0: default '\n '1: value embedding + temporal embedding + positional embedding '\n '2: value embedding + temporal embedding '\n '3: value embedding + positional embedding '\n '4: value embedding')\nparser.add_argument('--enc_in', type=int, default=7, help='encoder input size')\nparser.add_argument('--dec_in', type=int, default=7, help='decoder input size')\nparser.add_argument('--c_out', type=int, default=7, help='output size')\nparser.add_argument('--d_model', type=int, default=512, help='dimension of model')\nparser.add_argument('--n_heads', type=int, default=8, help='num of heads')\nparser.add_argument('--e_layers', type=int, default=2, help='num of encoder layers')\nparser.add_argument('--d_layers', type=int, default=1, help='num of decoder layers')\nparser.add_argument('--d_ff', type=int, default=2048, help='dimension of fcn')\nparser.add_argument('--moving_avg', type=int, default=25, help='window size of moving average')\nparser.add_argument('--factor', type=int, default=1, help='attn factor')\nparser.add_argument('--distil', action='store_false',\n help='whether to use distilling in encoder, using this argument means not using distilling',\n default=True)\nparser.add_argument('--dropout', type=float, default=0.05, help='dropout')\nparser.add_argument('--embed', type=str, default='timeF',\n help='time features encoding, options:[timeF, fixed, learned]')\nparser.add_argument('--activation', type=str, default='gelu', help='activation')\nparser.add_argument('--output_attention', action='store_true', help='whether to output attention in encoder')\nparser.add_argument('--do_predict', action='store_true', help='whether to predict unseen future data')\n\n# optimization\nparser.add_argument('--num_workers', type=int, default=8, help='data loader num workers')\nparser.add_argument('--itr', type=int, default=2, help='experiments times')\nparser.add_argument('--train_epochs', type=int, default=10, help='train epochs')\nparser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')\nparser.add_argument('--patience', type=int, default=3, help='early stopping patience')\nparser.add_argument('--learning_rate', type=float, default=0.0001, help='optimizer learning rate')\nparser.add_argument('--des', type=str, default='test', help='exp description')\nparser.add_argument('--loss', type=str, default='MSE', help='loss function')\nparser.add_argument('--lradj', type=str, default='type1', help='adjust learning rate')\nparser.add_argument('--use_amp', action='store_true', help='use automatic mixed precision training', default=False)\n\n# GPU\nparser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')\nparser.add_argument('--gpu', type=int, default=0, help='gpu')\nparser.add_argument('--use_multi_gpu', action='store_true', help='use multiple gpus', default=False)\nparser.add_argument('--devices', type=str, default='0,1,2,3', help='device ids of multile gpus')\nparser.add_argument('--test_flop', action='store_true', default=False, help='See utils/tools for usage')\n\nargs = parser.parse_args()\nargs.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False\n\nif args.use_gpu and args.use_multi_gpu:\n args.dvices = args.devices.replace(' ', '')\n device_ids = args.devices.split(',')\n args.device_ids = [int(id_) for id_ in device_ids]\n args.gpu = args.device_ids[0]\n\nprint('Args in experiment:')\nprint(args)\n\nExp = Exp_Main\n\nif args.is_training:\n start = time.time()\n for ii in range(args.itr):\n # setting record of experiments\n setting = '{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(\n args.model_id,\n args.model,\n args.data,\n args.features,\n args.seq_len,\n args.label_len,\n args.pred_len,\n args.d_model,\n args.n_heads,\n args.e_layers,\n args.d_layers,\n args.d_ff,\n args.factor,\n args.embed,\n args.distil,\n args.des, ii)\n\n exp = Exp(args) # set experiments\n print('>>>>>>>start training : {}>>>>>>>>>>>>>>>>>>>>>>>>>>'.format(setting))\n exp.train(setting)\n\n print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))\n exp.test(setting)\n\n # if args.do_predict:\n # print('>>>>>>>predicting : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))\n # exp.predict(setting, True)\n\n torch.cuda.empty_cache()\n end = time.time()\n used_time = end -start\n print(\"time:\",used_time)\n f = open(\"result.txt\", 'a')\n f.write('time:{}'.format(used_time))\n f.write('\\n')\n f.write('\\n')\n f.close()\nelse:\n ii = 0\n setting = '{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(args.model_id,\n args.model,\n args.data,\n args.features,\n args.seq_len,\n args.label_len,\n args.pred_len,\n args.d_model,\n args.n_heads,\n args.e_layers,\n args.d_layers,\n args.d_ff,\n args.factor,\n args.embed,\n args.distil,\n args.des, ii)\n\n exp = Exp(args) # set experiments\n print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))\n exp.test(setting, test=1)\n torch.cuda.empty_cache()\n" }, { "path": "scripts/ETTh1.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/ETTh1\" ]; then\n mkdir ./logs/ETTh1\nfi\nexport CUDA_VISIBLE_DEVICES=0\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh1.csv \\\n --model_id ETTh1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh1 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 64 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.1 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh1/$model_name'_'ETTh1_$seq_len'_'$pred_len.log\n\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh1.csv \\\n --model_id ETTh1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh1 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 64 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.1 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh1/$model_name'_'ETTh1_$seq_len'_'$pred_len.log\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh1.csv \\\n --model_id ETTh1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh1 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 64 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.1 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh1/$model_name'_'ETTh1_$seq_len'_'$pred_len.log\n\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh1.csv \\\n --model_id ETTh1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh1 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 16 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.1 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh1/$model_name'_'ETTh1_$seq_len'_'$pred_len.log" }, { "path": "scripts/ETTh2.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/ETTh2\" ]; then\n mkdir ./logs/ETTh2\nfi\nexport CUDA_VISIBLE_DEVICES=1\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh2.csv \\\n --model_id ETTh2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh2 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 16 \\\n --d_ff 32 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --top_k 5 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh2/$model_name'_'ETTh2_$seq_len'_'$pred_len.log\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh2.csv \\\n --model_id ETTh2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh2 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 16 \\\n --d_ff 32 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --top_k 5 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh2/$model_name'_'ETTh2_$seq_len'_'$pred_len.log\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh2.csv \\\n --model_id ETTh2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh2 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 16 \\\n --d_ff 32 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --top_k 5 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh2/$model_name'_'ETTh2_$seq_len'_'$pred_len.log\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTh2.csv \\\n --model_id ETTh2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTh2 \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 16 \\\n --d_ff 32 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --top_k 5 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTh2/$model_name'_'ETTh2_$seq_len'_'$pred_len.log" }, { "path": "scripts/ETTm1.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/ETTm1\" ]; then\n mkdir ./logs/ETTm1\nfi\n\nexport CUDA_VISIBLE_DEVICES=2\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm1.csv \\\n --model_id ETTm1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm1 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm1/$model_name'_'ETTm1_$seq_len'_'$pred_len.log\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm1.csv \\\n --model_id ETTm1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm1 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm1/$model_name'_'ETTm1_$seq_len'_'$pred_len.log\n\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm1.csv \\\n --model_id ETTm1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm1 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm1/$model_name'_'ETTm1_$seq_len'_'$pred_len.log\n\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm1.csv \\\n --model_id ETTm1'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm1 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm1/$model_name'_'ETTm1_$seq_len'_'$pred_len.log" }, { "path": "scripts/ETTm2.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/ETTm2\" ]; then\n mkdir ./logs/ETTm2\nfi\n\nexport CUDA_VISIBLE_DEVICES=3\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm2.csv \\\n --model_id ETTm2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm2 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.3 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm2/$model_name'_'ETTm2_$seq_len'_'$pred_len.log\n\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm2.csv \\\n --model_id ETTm2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm2 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 64 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.3 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm2/$model_name'_'ETTm2_$seq_len'_'$pred_len.log\n\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm2.csv \\\n --model_id ETTm2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm2 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 32 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.3 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm2/$model_name'_'ETTm2_$seq_len'_'$pred_len.log\n\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path ETTm2.csv \\\n --model_id ETTm2'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data ETTm2 \\\n --features M \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --d_model 32 \\\n --d_ff 64 \\\n --top_k 3 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --dropout 0.3 \\\n --batch_size 32 \\\n --itr 1 #>logs/ETTm2/$model_name'_'ETTm2_$seq_len'_'$pred_len.log\n" }, { "path": "scripts/Flight.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/Flight\" ]; then\n mkdir ./logs/Flight\nfi\n\nexport CUDA_VISIBLE_DEVICES=2\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\nfor pred_len in 96 192 336 720\ndo\n python -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path Flight.csv \\\n --model_id Flight'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'UUEE' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 7 \\\n --dec_in 7 \\\n --c_out 7 \\\n --des 'Exp' \\\n --itr 1 \\\n --d_model 16 \\\n --d_ff 32 \\\n --top_k 5 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --node_dim 100 \\\n --batch_size 32 #>logs/Flight/$model_name'_'Flight_$seq_len'_'$pred_len.log\n\ndone\n\n" }, { "path": "scripts/electricity.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/electricity\" ]; then\n mkdir ./logs/electricity\nfi\n\nexport CUDA_VISIBLE_DEVICES=3\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path electricity.csv \\\n --model_id electricity'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 321 \\\n --dec_in 321 \\\n --c_out 321 \\\n --des 'Exp' \\\n --d_model 1024 \\\n --d_ff 512 \\\n --top_k 5 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --node_dim 100 \\\n --batch_size 32 \\\n --itr 1 #>logs/electricity/$model_name'_'electricity_$seq_len'_'$pred_len.log\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path electricity.csv \\\n --model_id electricity'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 321 \\\n --dec_in 321 \\\n --c_out 321 \\\n --des 'Exp' \\\n --d_model 1024 \\\n --d_ff 512 \\\n --top_k 5 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --node_dim 100 \\\n --batch_size 32 \\\n --itr 1 #>logs/electricity/$model_name'_'electricity_$seq_len'_'$pred_len.log\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path electricity.csv \\\n --model_id electricity'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 3 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 321 \\\n --dec_in 321 \\\n --c_out 321 \\\n --des 'Exp' \\\n --d_model 1024 \\\n --d_ff 512 \\\n --top_k 5 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --node_dim 100 \\\n --batch_size 32 \\\n --itr 1 #>logs/electricity/$model_name'_'electricity_$seq_len'_'$pred_len.log\n\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path electricity.csv \\\n --model_id electricity'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 3 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 321 \\\n --dec_in 321 \\\n --c_out 321 \\\n --des 'Exp' \\\n --d_model 1024 \\\n --d_ff 512 \\\n --top_k 5 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --node_dim 100 \\\n --batch_size 32 \\\n --itr 1 #>logs/electricity/$model_name'_'electricity_$seq_len'_'$pred_len.log" }, { "path": "scripts/exchange.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/exchange\" ]; then\n mkdir ./logs/exchange\nfi\n\nexport CUDA_VISIBLE_DEVICES=2\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path exchange_rate.csv \\\n --model_id exchange'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 8 \\\n --dec_in 8 \\\n --c_out 8 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 3 \\\n --dropout 0.2 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/exchange/$model_name'_'exchange_$seq_len'_'$pred_len.log\n\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path exchange_rate.csv \\\n --model_id exchange'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 8 \\\n --dec_in 8 \\\n --c_out 8 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --node_dim 30 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/exchange/$model_name'_'exchange_$seq_len'_'$pred_len.log\n\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path exchange_rate.csv \\\n --model_id exchange'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 8 \\\n --dec_in 8 \\\n --c_out 8 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --node_dim 30 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/exchange/$model_name'_'exchange_$seq_len'_'$pred_len.log\n\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path exchange_rate.csv \\\n --model_id exchange'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 8 \\\n --dec_in 8 \\\n --c_out 8 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --conv_channel 16 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/exchange/$model_name'_'exchange_$seq_len'_'$pred_len.log" }, { "path": "scripts/weather.sh", "content": "if [ ! -d \"./logs\" ]; then\n mkdir ./logs\nfi\n\nif [ ! -d \"./logs/weather\" ]; then\n mkdir ./logs/weather\nfi\nexport CUDA_VISIBLE_DEVICES=2\n\nseq_len=96\nlabel_len=48\nmodel_name=MSGNet\n\npred_len=96\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path weather.csv \\\n --model_id weather'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 21 \\\n --dec_in 21 \\\n --c_out 21 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --train_epochs 3 \\\n --itr 1 #>logs/weather/$model_name'_'weather_$seq_len'_'$pred_len.log\n\npred_len=192\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path weather.csv \\\n --model_id weather'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 21 \\\n --dec_in 21 \\\n --c_out 21 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/weather/$model_name'_'weather_$seq_len'_'$pred_len.log\n\npred_len=336\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path weather.csv \\\n --model_id weather'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 1 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 21 \\\n --dec_in 21 \\\n --c_out 21 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/weather/$model_name'_'weather_$seq_len'_'$pred_len.log\n\npred_len=720\npython -u run_longExp.py \\\n --is_training 1 \\\n --root_path ./dataset/ \\\n --data_path weather.csv \\\n --model_id weather'_'$seq_len'_'$pred_len \\\n --model $model_name \\\n --data custom \\\n --features M \\\n --freq h \\\n --target 'OT' \\\n --seq_len $seq_len \\\n --label_len $label_len \\\n --pred_len $pred_len \\\n --e_layers 2 \\\n --d_layers 1 \\\n --factor 3 \\\n --enc_in 21 \\\n --dec_in 21 \\\n --c_out 21 \\\n --des 'Exp' \\\n --d_model 64 \\\n --d_ff 128 \\\n --top_k 5 \\\n --conv_channel 32 \\\n --skip_channel 32 \\\n --batch_size 32 \\\n --itr 1 #>logs/weather/$model_name'_'weather_$seq_len'_'$pred_len.log\n\n\n" }, { "path": "utils/masking.py", "content": "import torch\n\n\nclass TriangularCausalMask():\n def __init__(self, B, L, device=\"cpu\"):\n mask_shape = [B, 1, L, L]\n with torch.no_grad():\n self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)\n\n @property\n def mask(self):\n return self._mask\n\n\nclass ProbMask():\n def __init__(self, B, H, L, index, scores, device=\"cpu\"):\n _mask = torch.ones(L, scores.shape[-1], dtype=torch.bool).to(device).triu(1)\n _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])\n indicator = _mask_ex[torch.arange(B)[:, None, None],\n torch.arange(H)[None, :, None],\n index, :].to(device)\n self._mask = indicator.view(scores.shape).to(device)\n\n @property\n def mask(self):\n return self._mask\n" }, { "path": "utils/metrics.py", "content": "import numpy as np\n\ndef MAE(pred, true):\n return np.mean(np.abs(pred - true))\n\ndef MAPE(pred, true):\n return np.mean(np.abs((pred - true) / true))\n\ndef ND(pred, true):\n return np.mean(np.abs(true - pred)) / np.mean(np.abs(true))\n\ndef MSE(pred, true):\n return np.mean((pred - true) ** 2)\n\ndef RMSE(pred, true):\n return np.sqrt(MSE(pred, true))\n\ndef NRMSE(pred, true):\n return np.sqrt(np.mean(np.power((pred - true), 2))) / (np.mean(np.abs(true)))\n\ndef RSE(pred, true):\n return np.sqrt(np.sum((true - pred) ** 2)) / np.sqrt(np.sum((true - true.mean()) ** 2))\n\n\ndef CORR(pred, true):\n u = ((true - true.mean(0)) * (pred - pred.mean(0))).sum(0)\n d = np.sqrt(((true - true.mean(0)) ** 2 * (pred - pred.mean(0)) ** 2).sum(0))\n d += 1e-12\n return 0.01*(u / d).mean(-1)\n\n\ndef MSPE(pred, true):\n return np.mean(np.square((pred - true) / true))\n\n\ndef metric(pred, true):\n mae = MAE(pred, true)\n mse = MSE(pred, true)\n rmse = RMSE(pred, true)\n mape = MAPE(pred, true)\n mspe = MSPE(pred, true)\n rse = RSE(pred, true)\n corr = CORR(pred, true)\n nd = ND(pred,true)\n nrmse = NRMSE(pred,true)\n\n return mae, mse, rmse, mape, mspe, rse , corr, nd, nrmse\n\ndef metric2(pred, true):\n mae = MAE(pred, true)\n mse = MSE(pred, true)\n rmse = RMSE(pred, true)\n mape = MAPE(pred, true)\n mspe = MSPE(pred, true)\n rse = RSE(pred, true)\n nd = ND(pred,true)\n nrmse = NRMSE(pred,true)\n\n return mae, mse, rmse, mape, mspe, rse , nd, nrmse" }, { "path": "utils/timefeatures.py", "content": "from typing import List\n\nimport numpy as np\nimport pandas as pd\nfrom pandas.tseries import offsets\nfrom pandas.tseries.frequencies import to_offset\n\n\nclass TimeFeature:\n def __init__(self):\n pass\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n pass\n\n def __repr__(self):\n return self.__class__.__name__ + \"()\"\n\n\nclass SecondOfMinute(TimeFeature):\n \"\"\"Minute of hour encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return index.second / 59.0 - 0.5\n\n\nclass MinuteOfHour(TimeFeature):\n \"\"\"Minute of hour encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return index.minute / 59.0 - 0.5\n\n\nclass HourOfDay(TimeFeature):\n \"\"\"Hour of day encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return index.hour / 23.0 - 0.5\n\n\nclass DayOfWeek(TimeFeature):\n \"\"\"Hour of day encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return index.dayofweek / 6.0 - 0.5\n\n\nclass DayOfMonth(TimeFeature):\n \"\"\"Day of month encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return (index.day - 1) / 30.0 - 0.5\n\n\nclass DayOfYear(TimeFeature):\n \"\"\"Day of year encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return (index.dayofyear - 1) / 365.0 - 0.5\n\n\nclass MonthOfYear(TimeFeature):\n \"\"\"Month of year encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return (index.month - 1) / 11.0 - 0.5\n\n\nclass WeekOfYear(TimeFeature):\n \"\"\"Week of year encoded as value between [-0.5, 0.5]\"\"\"\n\n def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:\n return (index.isocalendar().week - 1) / 52.0 - 0.5\n\n\ndef time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:\n \"\"\"\n Returns a list of time features that will be appropriate for the given frequency string.\n Parameters\n ----------\n freq_str\n Frequency string of the form [multiple][granularity] such as \"12H\", \"5min\", \"1D\" etc.\n \"\"\"\n\n features_by_offsets = {\n offsets.YearEnd: [],\n offsets.QuarterEnd: [MonthOfYear],\n offsets.MonthEnd: [MonthOfYear],\n offsets.Week: [DayOfMonth, WeekOfYear],\n offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],\n offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],\n offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],\n offsets.Minute: [\n MinuteOfHour,\n HourOfDay,\n DayOfWeek,\n DayOfMonth,\n DayOfYear,\n ],\n offsets.Second: [\n SecondOfMinute,\n MinuteOfHour,\n HourOfDay,\n DayOfWeek,\n DayOfMonth,\n DayOfYear,\n ],\n }\n\n offset = to_offset(freq_str)\n\n for offset_type, feature_classes in features_by_offsets.items():\n if isinstance(offset, offset_type):\n return [cls() for cls in feature_classes]\n\n supported_freq_msg = f\"\"\"\n Unsupported frequency {freq_str}\n The following frequencies are supported:\n Y - yearly\n alias: A\n M - monthly\n W - weekly\n D - daily\n B - business days\n H - hourly\n T - minutely\n alias: min\n S - secondly\n \"\"\"\n raise RuntimeError(supported_freq_msg)\n\n\ndef time_features(dates, freq='h'):\n return np.vstack([feat(dates) for feat in time_features_from_frequency_str(freq)])\n" }, { "path": "utils/tools.py", "content": "import numpy as np\nimport torch\nimport matplotlib.pyplot as plt\nimport time\n\nplt.switch_backend('agg')\n\n\ndef adjust_learning_rate(optimizer, epoch, args):\n # lr = args.learning_rate * (0.2 ** (epoch // 2))\n if args.lradj == 'type1':\n lr_adjust = {epoch: args.learning_rate * (0.5 ** ((epoch - 1) // 1))}\n elif args.lradj == 'type2':\n lr_adjust = {\n 2: 5e-5, 4: 1e-5, 6: 5e-6, 8: 1e-6,\n 10: 5e-7, 15: 1e-7, 20: 5e-8\n }\n elif args.lradj == '3':\n lr_adjust = {epoch: args.learning_rate if epoch < 10 else args.learning_rate*0.1}\n elif args.lradj == '4':\n lr_adjust = {epoch: args.learning_rate if epoch < 15 else args.learning_rate*0.1}\n elif args.lradj == '5':\n lr_adjust = {epoch: args.learning_rate if epoch < 25 else args.learning_rate*0.1}\n elif args.lradj == '6':\n lr_adjust = {epoch: args.learning_rate if epoch < 5 else args.learning_rate*0.1} \n if epoch in lr_adjust.keys():\n lr = lr_adjust[epoch]\n for param_group in optimizer.param_groups:\n param_group['lr'] = lr\n print('Updating learning rate to {}'.format(lr))\n\n\nclass EarlyStopping:\n def __init__(self, patience=7, verbose=False, delta=0):\n self.patience = patience\n self.verbose = verbose\n self.counter = 0\n self.best_score = None\n self.early_stop = False\n self.val_loss_min = np.Inf\n self.delta = delta\n\n def __call__(self, val_loss, model, path):\n score = -val_loss\n if self.best_score is None:\n self.best_score = score\n self.save_checkpoint(val_loss, model, path)\n elif score < self.best_score + self.delta:\n self.counter += 1\n print(f'EarlyStopping counter: {self.counter} out of {self.patience}')\n if self.counter >= self.patience:\n self.early_stop = True\n else:\n self.best_score = score\n self.save_checkpoint(val_loss, model, path)\n self.counter = 0\n\n def save_checkpoint(self, val_loss, model, path):\n if self.verbose:\n print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...')\n torch.save(model.state_dict(), path + '/' + 'checkpoint.pth')\n self.val_loss_min = val_loss\n\n\nclass dotdict(dict):\n \"\"\"dot.notation access to dictionary attributes\"\"\"\n __getattr__ = dict.get\n __setattr__ = dict.__setitem__\n __delattr__ = dict.__delitem__\n\n\nclass StandardScaler():\n def __init__(self, mean, std):\n self.mean = mean\n self.std = std\n\n def transform(self, data):\n return (data - self.mean) / self.std\n\n def inverse_transform(self, data):\n return (data * self.std) + self.mean\n\n\ndef visual(true, preds=None, name='./pic/test.pdf'):\n \"\"\"\n Results visualization\n \"\"\"\n plt.figure()\n plt.plot(true, label='GroundTruth', linewidth=2)\n if preds is not None:\n plt.plot(preds, label='Prediction', linewidth=2)\n plt.legend()\n plt.show()\n plt.savefig(name, bbox_inches='tight')\n\ndef test_params_flop(model,x_shape):\n \"\"\"\n If you want to thest former's flop, you need to give default value to inputs in model.forward(), the following code can only pass one argument to forward()\n \"\"\"\n model_params = 0\n for parameter in model.parameters():\n model_params += parameter.numel()\n print('INFO: Trainable parameter count: {:.2f}M'.format(model_params / 1000000.0))\n from ptflops import get_model_complexity_info \n with torch.cuda.device(0):\n macs, params = get_model_complexity_info(model.cuda(), x_shape, as_strings=True, print_per_layer_stat=True)\n # print('Flops:' + flops)\n # print('Params:' + params)\n print('{:<30} {:<8}'.format('Computational complexity: ', macs))\n print('{:<30} {:<8}'.format('Number of parameters: ', params))" } ]