add STAEformer to BasicTS (#80)

GestaltCogTeam · Oct 26, 2023 · 2531b2a · 2531b2a
1 parent 6a49da5
commit 2531b2a
Show file tree

Hide file tree

Showing 10 changed files with 910 additions and 1 deletion.
diff --git a/baselines/STAEformer/STAEformer_METR-LA.py b/baselines/STAEformer/STAEformer_METR-LA.py
@@ -0,0 +1,108 @@
+import os
+import sys
+
+# TODO: remove it when basicts can be installed by pip
+sys.path.append(os.path.abspath(__file__ + "/../../.."))
+import torch
+from easydict import EasyDict
+from basicts.archs import STAEformer
+from basicts.runners import SimpleTimeSeriesForecastingRunner
+from basicts.data import TimeSeriesForecastingDataset
+from basicts.losses import masked_mae
+from basicts.utils import load_adj
+
+
+CFG = EasyDict()
+
+# ================= general ================= #
+CFG.DESCRIPTION = "STAEformer model configuration"
+CFG.RUNNER = SimpleTimeSeriesForecastingRunner
+CFG.DATASET_CLS = TimeSeriesForecastingDataset
+CFG.DATASET_NAME = "METR-LA"
+CFG.DATASET_TYPE = "Traffic speed"
+CFG.DATASET_INPUT_LEN = 12
+CFG.DATASET_OUTPUT_LEN = 12
+CFG.GPU_NUM = 1
+CFG.NULL_VAL = 0.0
+
+# ================= environment ================= #
+CFG.ENV = EasyDict()
+CFG.ENV.SEED = 1
+CFG.ENV.CUDNN = EasyDict()
+CFG.ENV.CUDNN.ENABLED = True
+
+# ================= model ================= #
+CFG.MODEL = EasyDict()
+CFG.MODEL.NAME = "STAEformer"
+CFG.MODEL.ARCH = STAEformer
+adj_mx, _ = load_adj("datasets/" + CFG.DATASET_NAME + "/adj_mx.pkl", "normlap")
+adj_mx = torch.Tensor(adj_mx[0])
+CFG.MODEL.PARAM = {
+    "num_nodes" : 207
+}
+CFG.MODEL.FORWARD_FEATURES = [0,1,2]
+CFG.MODEL.TARGET_FEATURES = [0]
+
+# ================= optim ================= #
+CFG.TRAIN = EasyDict()
+CFG.TRAIN.LOSS = masked_mae
+CFG.TRAIN.OPTIM = EasyDict()
+CFG.TRAIN.OPTIM.TYPE = "Adam"
+CFG.TRAIN.OPTIM.PARAM = {
+    "lr": 0.001,
+    "weight_decay": 0.0003,
+}
+CFG.TRAIN.LR_SCHEDULER = EasyDict()
+CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
+CFG.TRAIN.LR_SCHEDULER.PARAM = {
+    "milestones": [20, 25],
+    "gamma": 0.1
+}
+
+# ================= train ================= #
+# CFG.TRAIN.CLIP_GRAD_PARAM = {
+#     "max_norm": 5.0
+# }
+CFG.TRAIN.NUM_EPOCHS = 30
+CFG.TRAIN.CKPT_SAVE_DIR = os.path.join(
+    "checkpoints",
+    "_".join([CFG.MODEL.NAME, str(CFG.TRAIN.NUM_EPOCHS)])
+)
+# train data
+CFG.TRAIN.DATA = EasyDict()
+# read data
+CFG.TRAIN.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TRAIN.DATA.BATCH_SIZE = 16
+CFG.TRAIN.DATA.PREFETCH = False
+CFG.TRAIN.DATA.SHUFFLE = True
+CFG.TRAIN.DATA.NUM_WORKERS = 2
+CFG.TRAIN.DATA.PIN_MEMORY = False
+
+# ================= validate ================= #
+CFG.VAL = EasyDict()
+CFG.VAL.INTERVAL = 1
+# validating data
+CFG.VAL.DATA = EasyDict()
+# read data
+CFG.VAL.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.VAL.DATA.BATCH_SIZE = 16
+CFG.VAL.DATA.PREFETCH = False
+CFG.VAL.DATA.SHUFFLE = False
+CFG.VAL.DATA.NUM_WORKERS = 2
+CFG.VAL.DATA.PIN_MEMORY = False
+
+# ================= test ================= #
+CFG.TEST = EasyDict()
+CFG.TEST.INTERVAL = 1
+# test data
+CFG.TEST.DATA = EasyDict()
+# read data
+CFG.TEST.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TEST.DATA.BATCH_SIZE = 16
+CFG.TEST.DATA.PREFETCH = False
+CFG.TEST.DATA.SHUFFLE = False
+CFG.TEST.DATA.NUM_WORKERS = 2
+CFG.TEST.DATA.PIN_MEMORY = False
diff --git a/baselines/STAEformer/STAEformer_PEMS-BAY.py b/baselines/STAEformer/STAEformer_PEMS-BAY.py
@@ -0,0 +1,108 @@
+import os
+import sys
+
+# TODO: remove it when basicts can be installed by pip
+sys.path.append(os.path.abspath(__file__ + "/../../.."))
+import torch
+from easydict import EasyDict
+from basicts.archs import STAEformer
+from basicts.runners import SimpleTimeSeriesForecastingRunner
+from basicts.data import TimeSeriesForecastingDataset
+from basicts.losses import masked_mae
+from basicts.utils import load_adj
+
+
+CFG = EasyDict()
+
+# ================= general ================= #
+CFG.DESCRIPTION = "STAEformer model configuration"
+CFG.RUNNER = SimpleTimeSeriesForecastingRunner
+CFG.DATASET_CLS = TimeSeriesForecastingDataset
+CFG.DATASET_NAME = "PEMS-BAY"
+CFG.DATASET_TYPE = "Traffic speed"
+CFG.DATASET_INPUT_LEN = 12
+CFG.DATASET_OUTPUT_LEN = 12
+CFG.GPU_NUM = 1
+CFG.NULL_VAL = 0.0
+
+# ================= environment ================= #
+CFG.ENV = EasyDict()
+CFG.ENV.SEED = 1
+CFG.ENV.CUDNN = EasyDict()
+CFG.ENV.CUDNN.ENABLED = True
+
+# ================= model ================= #
+CFG.MODEL = EasyDict()
+CFG.MODEL.NAME = "STAEformer"
+CFG.MODEL.ARCH = STAEformer
+adj_mx, _ = load_adj("datasets/" + CFG.DATASET_NAME + "/adj_mx.pkl", "normlap")
+adj_mx = torch.Tensor(adj_mx[0])
+CFG.MODEL.PARAM = {
+    "num_nodes" : 207
+}
+CFG.MODEL.FORWARD_FEATURES = [0,1,2]
+CFG.MODEL.TARGET_FEATURES = [0]
+
+# ================= optim ================= #
+CFG.TRAIN = EasyDict()
+CFG.TRAIN.LOSS = masked_mae
+CFG.TRAIN.OPTIM = EasyDict()
+CFG.TRAIN.OPTIM.TYPE = "Adam"
+CFG.TRAIN.OPTIM.PARAM = {
+    "lr": 0.001,
+    "weight_decay": 0.0003,
+}
+CFG.TRAIN.LR_SCHEDULER = EasyDict()
+CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
+CFG.TRAIN.LR_SCHEDULER.PARAM = {
+    "milestones": [20, 25],
+    "gamma": 0.1
+}
+
+# ================= train ================= #
+# CFG.TRAIN.CLIP_GRAD_PARAM = {
+#     "max_norm": 5.0
+# }
+CFG.TRAIN.NUM_EPOCHS = 30
+CFG.TRAIN.CKPT_SAVE_DIR = os.path.join(
+    "checkpoints",
+    "_".join([CFG.MODEL.NAME, str(CFG.TRAIN.NUM_EPOCHS)])
+)
+# train data
+CFG.TRAIN.DATA = EasyDict()
+# read data
+CFG.TRAIN.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TRAIN.DATA.BATCH_SIZE = 16
+CFG.TRAIN.DATA.PREFETCH = False
+CFG.TRAIN.DATA.SHUFFLE = True
+CFG.TRAIN.DATA.NUM_WORKERS = 2
+CFG.TRAIN.DATA.PIN_MEMORY = False
+
+# ================= validate ================= #
+CFG.VAL = EasyDict()
+CFG.VAL.INTERVAL = 1
+# validating data
+CFG.VAL.DATA = EasyDict()
+# read data
+CFG.VAL.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.VAL.DATA.BATCH_SIZE = 16
+CFG.VAL.DATA.PREFETCH = False
+CFG.VAL.DATA.SHUFFLE = False
+CFG.VAL.DATA.NUM_WORKERS = 2
+CFG.VAL.DATA.PIN_MEMORY = False
+
+# ================= test ================= #
+CFG.TEST = EasyDict()
+CFG.TEST.INTERVAL = 1
+# test data
+CFG.TEST.DATA = EasyDict()
+# read data
+CFG.TEST.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TEST.DATA.BATCH_SIZE = 16
+CFG.TEST.DATA.PREFETCH = False
+CFG.TEST.DATA.SHUFFLE = False
+CFG.TEST.DATA.NUM_WORKERS = 2
+CFG.TEST.DATA.PIN_MEMORY = False
diff --git a/baselines/STAEformer/STAEformer_PEMS03.py b/baselines/STAEformer/STAEformer_PEMS03.py
@@ -0,0 +1,108 @@
+import os
+import sys
+
+# TODO: remove it when basicts can be installed by pip
+sys.path.append(os.path.abspath(__file__ + "/../../.."))
+import torch
+from easydict import EasyDict
+from basicts.archs import STAEformer
+from basicts.runners import SimpleTimeSeriesForecastingRunner
+from basicts.data import TimeSeriesForecastingDataset
+from basicts.losses import masked_mae
+from basicts.utils import load_adj
+
+
+CFG = EasyDict()
+
+# ================= general ================= #
+CFG.DESCRIPTION = "STAEformer model configuration"
+CFG.RUNNER = SimpleTimeSeriesForecastingRunner
+CFG.DATASET_CLS = TimeSeriesForecastingDataset
+CFG.DATASET_NAME = "PEMS03"
+CFG.DATASET_TYPE = "Traffic flow"
+CFG.DATASET_INPUT_LEN = 12
+CFG.DATASET_OUTPUT_LEN = 12
+CFG.GPU_NUM = 1
+CFG.NULL_VAL = 0.0
+
+# ================= environment ================= #
+CFG.ENV = EasyDict()
+CFG.ENV.SEED = 1
+CFG.ENV.CUDNN = EasyDict()
+CFG.ENV.CUDNN.ENABLED = True
+
+# ================= model ================= #
+CFG.MODEL = EasyDict()
+CFG.MODEL.NAME = "STAEformer"
+CFG.MODEL.ARCH = STAEformer
+adj_mx, _ = load_adj("datasets/" + CFG.DATASET_NAME + "/adj_mx.pkl", "normlap")
+adj_mx = torch.Tensor(adj_mx[0])
+CFG.MODEL.PARAM = {
+    "num_nodes" : 170
+}
+CFG.MODEL.FORWARD_FEATURES = [0,1,2]
+CFG.MODEL.TARGET_FEATURES = [0]
+
+# ================= optim ================= #
+CFG.TRAIN = EasyDict()
+CFG.TRAIN.LOSS = masked_mae
+CFG.TRAIN.OPTIM = EasyDict()
+CFG.TRAIN.OPTIM.TYPE = "Adam"
+CFG.TRAIN.OPTIM.PARAM = {
+    "lr": 0.001,
+    "weight_decay": 0.0015,
+}
+CFG.TRAIN.LR_SCHEDULER = EasyDict()
+CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
+CFG.TRAIN.LR_SCHEDULER.PARAM = {
+    "milestones": [25, 45, 65],
+    "gamma": 0.1
+}
+
+# ================= train ================= #
+# CFG.TRAIN.CLIP_GRAD_PARAM = {
+#     "max_norm": 5.0
+# }
+CFG.TRAIN.NUM_EPOCHS = 70
+CFG.TRAIN.CKPT_SAVE_DIR = os.path.join(
+    "checkpoints",
+    "_".join([CFG.MODEL.NAME, str(CFG.TRAIN.NUM_EPOCHS)])
+)
+# train data
+CFG.TRAIN.DATA = EasyDict()
+# read data
+CFG.TRAIN.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TRAIN.DATA.BATCH_SIZE = 16
+CFG.TRAIN.DATA.PREFETCH = False
+CFG.TRAIN.DATA.SHUFFLE = True
+CFG.TRAIN.DATA.NUM_WORKERS = 2
+CFG.TRAIN.DATA.PIN_MEMORY = False
+
+# ================= validate ================= #
+CFG.VAL = EasyDict()
+CFG.VAL.INTERVAL = 1
+# validating data
+CFG.VAL.DATA = EasyDict()
+# read data
+CFG.VAL.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.VAL.DATA.BATCH_SIZE = 16
+CFG.VAL.DATA.PREFETCH = False
+CFG.VAL.DATA.SHUFFLE = False
+CFG.VAL.DATA.NUM_WORKERS = 2
+CFG.VAL.DATA.PIN_MEMORY = False
+
+# ================= test ================= #
+CFG.TEST = EasyDict()
+CFG.TEST.INTERVAL = 1
+# test data
+CFG.TEST.DATA = EasyDict()
+# read data
+CFG.TEST.DATA.DIR = "datasets/" + CFG.DATASET_NAME
+# dataloader args, optional
+CFG.TEST.DATA.BATCH_SIZE = 16
+CFG.TEST.DATA.PREFETCH = False
+CFG.TEST.DATA.SHUFFLE = False
+CFG.TEST.DATA.NUM_WORKERS = 2
+CFG.TEST.DATA.PIN_MEMORY = False