add a new ST baseline GMSDR

README.md

@@ -69,7 +69,7 @@ BasicTS support a variety of datasets, including spatial-temporal forecasting, l
 ### Baselines
 
 BasicTS provides a wealth of built-in models, including both spatial-temporal forecasting models and long time-series forecasting models, e.g.,
-- DCRNN, Graph WaveNet, MTGNN, STID, D2STGNN, STEP, DGCRN, DGCRN, STNorm, AGCRN, GTS, StemGNN, MegaCRN, STGCN, ...
+- DCRNN, Graph WaveNet, MTGNN, STID, D2STGNN, STEP, DGCRN, DGCRN, STNorm, AGCRN, GTS, StemGNN, MegaCRN, STGCN, STWave, STAEformer, GMSDR, ...
 - Informer, Autoformer, FEDformer, Pyraformer, DLinear, NLinear, Triformer, Crossformer, ...
 
 ## 💿 Dependencies

baselines/GMSDR/METR-LA.py

@@ -0,0 +1,130 @@
+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.runners import SimpleTimeSeriesForecastingRunner
+from basicts.data import TimeSeriesForecastingDataset
+from basicts.losses import masked_mae
+from basicts.utils import load_adj
+
+from .arch import GMSDR
+
+CFG = EasyDict()
+
+# DCRNN does not allow to load parameters since it creates parameters in the first iteration
+resume = False
+if not resume:
+    import random
+    _ = random.randint(-1e6, 1e6)
+
+# ================= general ================= #
+CFG.DESCRIPTION = "GMSDR 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._ = _
+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 = "GMSDR"
+CFG.MODEL.ARCH = GMSDR
+adj_mx, _ = load_adj("datasets/" + CFG.DATASET_NAME +
+                     "/adj_mx.pkl", "doubletransition")
+CFG.MODEL.PARAM = {
+    "horizon": 12,
+    "input_dim": 1,
+    "max_diffusion_step": 1,
+    "num_nodes": 207,
+    "num_rnn_layers": 2,
+    "output_dim": 1,
+    "rnn_units": 64,
+    "seq_len": 12,
+    "adj_mx": [torch.tensor(i) for i in adj_mx],
+    "pre_k": 6,
+    "pre_v": 1,
+}
+CFG.MODEL.FORWARD_FEATURES = [0]
+CFG.MODEL.TARGET_FEATURES = [0]
+CFG.MODEL.SETUP_GRAPH = True
+
+# ================= optim ================= #
+CFG.TRAIN = EasyDict()
+CFG.TRAIN.LOSS = masked_mae
+CFG.TRAIN.OPTIM = EasyDict()
+CFG.TRAIN.OPTIM.TYPE = "Adam"
+CFG.TRAIN.OPTIM.PARAM = {
+    "lr": 0.0015,
+    "eps": 1e-3
+}
+CFG.TRAIN.LR_SCHEDULER = EasyDict()
+CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
+CFG.TRAIN.LR_SCHEDULER.PARAM = {
+    "milestones": [30, 50, 70, 80],
+    "gamma": 0.2
+}
+
+# ================= train ================= #
+CFG.TRAIN.CLIP_GRAD_PARAM = {
+    "max_norm": 5.0
+}
+CFG.TRAIN.NUM_EPOCHS = 100
+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 = 64
+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 = 64
+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 = 64
+CFG.TEST.DATA.PREFETCH = False
+CFG.TEST.DATA.SHUFFLE = False
+CFG.TEST.DATA.NUM_WORKERS = 2
+CFG.TEST.DATA.PIN_MEMORY = False
+
+# ================= evaluate ================= #
+CFG.EVAL = EasyDict()
+CFG.EVAL.HORIZONS = [3, 6, 12]

baselines/GMSDR/PEMS-BAY.py

@@ -0,0 +1,130 @@
+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.runners import SimpleTimeSeriesForecastingRunner
+from basicts.data import TimeSeriesForecastingDataset
+from basicts.losses import masked_mae
+from basicts.utils import load_adj
+
+from .arch import GMSDR
+
+CFG = EasyDict()
+
+# DCRNN does not allow to load parameters since it creates parameters in the first iteration
+resume = False
+if not resume:
+    import random
+    _ = random.randint(-1e6, 1e6)
+
+# ================= general ================= #
+CFG.DESCRIPTION = "GMSDR 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._ = _
+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 = "GMSDR"
+CFG.MODEL.ARCH = GMSDR
+adj_mx, _ = load_adj("datasets/" + CFG.DATASET_NAME +
+                     "/adj_mx.pkl", "doubletransition")
+CFG.MODEL.PARAM = {
+    "horizon": 12,
+    "input_dim": 1,
+    "max_diffusion_step": 1,
+    "num_nodes": 325,
+    "num_rnn_layers": 2,
+    "output_dim": 1,
+    "rnn_units": 64,
+    "seq_len": 12,
+    "adj_mx": [torch.tensor(i) for i in adj_mx],
+    "pre_k": 6,
+    "pre_v": 1,
+}
+CFG.MODEL.FORWARD_FEATURES = [0]
+CFG.MODEL.TARGET_FEATURES = [0]
+CFG.MODEL.SETUP_GRAPH = True
+
+# ================= optim ================= #
+CFG.TRAIN = EasyDict()
+CFG.TRAIN.LOSS = masked_mae
+CFG.TRAIN.OPTIM = EasyDict()
+CFG.TRAIN.OPTIM.TYPE = "Adam"
+CFG.TRAIN.OPTIM.PARAM = {
+    "lr": 0.0015,
+    "eps": 1e-3
+}
+CFG.TRAIN.LR_SCHEDULER = EasyDict()
+CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
+CFG.TRAIN.LR_SCHEDULER.PARAM = {
+    "milestones": [30, 50, 70, 80],
+    "gamma": 0.2
+}
+
+# ================= train ================= #
+CFG.TRAIN.CLIP_GRAD_PARAM = {
+    "max_norm": 5.0
+}
+CFG.TRAIN.NUM_EPOCHS = 100
+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 = 64
+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 = 64
+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 = 64
+CFG.TEST.DATA.PREFETCH = False
+CFG.TEST.DATA.SHUFFLE = False
+CFG.TEST.DATA.NUM_WORKERS = 2
+CFG.TEST.DATA.PIN_MEMORY = False
+
+# ================= evaluate ================= #
+CFG.EVAL = EasyDict()
+CFG.EVAL.HORIZONS = [3, 6, 12]