diff --git a/README.md b/README.md index bfabefbb..61199a92 100644 --- a/README.md +++ b/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 diff --git a/baselines/GMSDR/METR-LA.py b/baselines/GMSDR/METR-LA.py new file mode 100644 index 00000000..468a5203 --- /dev/null +++ b/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] diff --git a/baselines/GMSDR/PEMS-BAY.py b/baselines/GMSDR/PEMS-BAY.py new file mode 100644 index 00000000..107750f6 --- /dev/null +++ b/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] diff --git a/baselines/GMSDR/PEMS03.py b/baselines/GMSDR/PEMS03.py new file mode 100644 index 00000000..2e21be03 --- /dev/null +++ b/baselines/GMSDR/PEMS03.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 = "PEMS03" +CFG.DATASET_TYPE = "Traffic flow" +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": 358, + "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] diff --git a/baselines/GMSDR/PEMS04.py b/baselines/GMSDR/PEMS04.py new file mode 100644 index 00000000..43f2a010 --- /dev/null +++ b/baselines/GMSDR/PEMS04.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 = "PEMS04" +CFG.DATASET_TYPE = "Traffic flow" +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": 307, + "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] diff --git a/baselines/GMSDR/PEMS07.py b/baselines/GMSDR/PEMS07.py new file mode 100644 index 00000000..2249225a --- /dev/null +++ b/baselines/GMSDR/PEMS07.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 = "PEMS07" +CFG.DATASET_TYPE = "Traffic flow" +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": 883, + "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] diff --git a/baselines/GMSDR/PEMS08.py b/baselines/GMSDR/PEMS08.py new file mode 100644 index 00000000..edd34231 --- /dev/null +++ b/baselines/GMSDR/PEMS08.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 = "PEMS08" +CFG.DATASET_TYPE = "Traffic flow" +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": 170, + "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] diff --git a/baselines/GMSDR/arch/__init__.py b/baselines/GMSDR/arch/__init__.py new file mode 100644 index 00000000..a56c811c --- /dev/null +++ b/baselines/GMSDR/arch/__init__.py @@ -0,0 +1,3 @@ +from .gmsdr_arch import GMSDR + +__all__ = ['GMSDR'] \ No newline at end of file diff --git a/baselines/GMSDR/arch/gmsdr_arch.py b/baselines/GMSDR/arch/gmsdr_arch.py new file mode 100644 index 00000000..859f5210 --- /dev/null +++ b/baselines/GMSDR/arch/gmsdr_arch.py @@ -0,0 +1,159 @@ +import numpy as np +import torch +import torch.nn as nn + +from .gmsdr_cell import GMSDRCell + +def count_parameters(model): + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +class Seq2SeqAttrs: + def __init__(self, adj_mx, **model_kwargs): + self.adj_mx = adj_mx + self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2)) + self.num_nodes = int(model_kwargs.get('num_nodes', 1)) + self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1)) + self.rnn_units = int(model_kwargs.get('rnn_units')) + self.hidden_state_size = self.num_nodes * self.rnn_units + self.pre_k = int(model_kwargs.get('pre_k', 1)) + self.pre_v = int(model_kwargs.get('pre_v', 1)) + self.input_dim = int(model_kwargs.get('input_dim', 1)) + self.output_dim = int(model_kwargs.get('output_dim', 1)) + + +class EncoderModel(nn.Module, Seq2SeqAttrs): + def __init__(self, adj_mx, **model_kwargs): + nn.Module.__init__(self) + Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) + self.input_dim = int(model_kwargs.get('input_dim', 1)) + self.seq_len = int(model_kwargs.get('seq_len')) # for the encoder + self.mlp = nn.Linear(self.input_dim, self.rnn_units) + self.gmsdr_layers = nn.ModuleList( + [GMSDRCell(self.rnn_units, self.input_dim, adj_mx, self.max_diffusion_step, self.num_nodes, self.pre_k, self.pre_v) for _ in range(self.num_rnn_layers)]) + + def forward(self, inputs, hx_k): + """ + Encoder forward pass. + + :param inputs: shape (batch_size, self.num_nodes * self.input_dim) + :param hx_k: (num_layers, batch_size, pre_k, self.num_nodes, self.rnn_units) + optional, zeros if not provided + :return: output: # shape (batch_size, self.hidden_state_size) + hx_k # shape (num_layers, batch_size, pre_k, self.num_nodes, self.rnn_units) + (lower indices mean lower layers) + """ + hx_ks = [] + batch = inputs.shape[0] + x = inputs.reshape(batch, self.num_nodes, self.input_dim) + output = self.mlp(x).view(batch, -1) + for layer_num, dcgru_layer in enumerate(self.gmsdr_layers): + next_hidden_state, new_hx_k = dcgru_layer(output, hx_k[layer_num]) + hx_ks.append(new_hx_k) + output = next_hidden_state + return output, torch.stack(hx_ks) + + +class DecoderModel(nn.Module, Seq2SeqAttrs): + def __init__(self, adj_mx, **model_kwargs): + nn.Module.__init__(self) + Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) + self.output_dim = int(model_kwargs.get('output_dim', 1)) + self.horizon = int(model_kwargs.get('horizon', 12)) # for the decoder + self.projection_layer = nn.Linear(self.rnn_units, self.output_dim) + self.gmsdr_layers = nn.ModuleList( + [GMSDRCell(self.rnn_units, self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes, self.pre_k, self.pre_v + ) for _ in range(self.num_rnn_layers)]) + + def forward(self, inputs, hx_k): + """ + Decoder forward pass. + + :param inputs: shape (batch_size, self.num_nodes * self.output_dim) + :param hx_k: (num_layers, batch_size, pre_k, num_nodes, rnn_units) + optional, zeros if not provided + :return: output: # shape (batch_size, self.num_nodes * self.output_dim) + hidden_state # shape (num_layers, batch_size, self.hidden_state_size) + (lower indices mean lower layers) + """ + hx_ks = [] + output = inputs + for layer_num, dcgru_layer in enumerate(self.gmsdr_layers): + next_hidden_state, new_hx_k = dcgru_layer(output, hx_k[layer_num]) + hx_ks.append(new_hx_k) + output = next_hidden_state + + projected = self.projection_layer(output.view(-1, self.rnn_units)) + output = projected.view(-1, self.num_nodes * self.output_dim) + + return output, torch.stack(hx_ks) + + +class GMSDR(nn.Module, Seq2SeqAttrs): + def __init__(self, adj_mx, **model_kwargs): + super().__init__() + Seq2SeqAttrs.__init__(self, adj_mx, **model_kwargs) + self.encoder_model = EncoderModel(adj_mx, **model_kwargs) + self.decoder_model = DecoderModel(adj_mx, **model_kwargs) + self.out = nn.Linear(self.rnn_units, self.decoder_model.output_dim) + + def encoder(self, inputs): + """ + encoder forward pass on t time steps + :param inputs: shape (seq_len, batch_size, num_sensor * input_dim) + :return: hx_k: (num_layers, batch_size, pre_k, num_sensor, rnn_units) + """ + hx_k = torch.zeros(self.num_rnn_layers, inputs.shape[1], self.pre_k, self.num_nodes, self.rnn_units, + device=inputs.device) + outputs = [] + for t in range(self.encoder_model.seq_len): + output, hx_k = self.encoder_model(inputs[t], hx_k) + outputs.append(output) + return torch.stack(outputs), hx_k + + def decoder(self, inputs, hx_k, labels=None, batches_seen=None): + """ + Decoder forward pass + :param inputs: (seq_len, batch_size, num_sensor * rnn_units) + :param hx_k: (num_layers, batch_size, pre_k, num_sensor, rnn_units) + :param labels: (self.horizon, batch_size, self.num_nodes * self.output_dim) [optional, not exist for inference] + :param batches_seen: global step [optional, not exist for inference] + :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim) + """ + decoder_hx_k = hx_k + decoder_input = inputs + + outputs = [] + for t in range(self.decoder_model.horizon): + decoder_output, decoder_hx_k = self.decoder_model(decoder_input[t], + decoder_hx_k) + outputs.append(decoder_output) + outputs = torch.stack(outputs) + return outputs + + def Loss_l2(self): + base_params = dict(self.named_parameters()) + loss_l2 = 0 + count = 0 + for key,value in base_params.items(): + if 'bias' not in key: + loss_l2 += torch.sum(value**2) + count += value.nelement() + return loss_l2 + + def forward(self, history_data, future_data=None, batch_seen=None, **kwargs): + """ + seq2seq forward pass + :param inputs: shape (seq_len, batch_size, num_sensor * input_dim) + :param labels: shape (horizon, batch_size, num_sensor * output) + :param batches_seen: batches seen till now + :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim) + """ + inputs = history_data.transpose(0,1).reshape(history_data.shape[1],history_data.shape[0],-1) + encoder_outputs, hx_k = self.encoder(inputs) + outputs = self.decoder(encoder_outputs, hx_k, future_data, batches_seen=batch_seen) + if batch_seen == 0: + print( + "Total trainable parameters {}".format(count_parameters(self)) + ) + return outputs.transpose(0,1).reshape(history_data.shape[0],history_data.shape[1],history_data.shape[2],-1) \ No newline at end of file diff --git a/baselines/GMSDR/arch/gmsdr_cell.py b/baselines/GMSDR/arch/gmsdr_cell.py new file mode 100644 index 00000000..0926a495 --- /dev/null +++ b/baselines/GMSDR/arch/gmsdr_cell.py @@ -0,0 +1,184 @@ +import numpy as np +import torch +from torch import nn, Tensor +import torch.nn.functional as F + + +class Seq2SeqAttrs: + def __init__(self, adj_mx, **model_kwargs): + self.adj_mx = adj_mx + self.max_diffusion_step = int(model_kwargs.get('max_diffusion_step', 2)) + self.num_nodes = int(model_kwargs.get('num_nodes', 1)) + self.num_rnn_layers = int(model_kwargs.get('num_rnn_layers', 1)) + self.rnn_units = int(model_kwargs.get('rnn_units')) + self.hidden_state_size = self.num_nodes * self.rnn_units + self.pre_k = int(model_kwargs.get('pre_k', 1)) + self.pre_v = int(model_kwargs.get('pre_v', 1)) + self.input_dim = int(model_kwargs.get('input_dim', 2)) + self.output_dim = int(model_kwargs.get('output_dim', 1)) + + +class LayerParams: + def __init__(self, rnn_network: torch.nn.Module, layer_type: str): + self._rnn_network = rnn_network + self._params_dict = {} + self._biases_dict = {} + self._type = layer_type + + def get_weights(self, shape): + if shape not in self._params_dict: + nn_param = torch.nn.Parameter(torch.empty(*shape)) + torch.nn.init.xavier_normal_(nn_param) + self._params_dict[shape] = nn_param + self._rnn_network.register_parameter('{}_weight_{}'.format(self._type, str(shape)), + nn_param) + return self._params_dict[shape] + + def get_biases(self, length, bias_start=0.0): + if length not in self._biases_dict: + biases = torch.nn.Parameter(torch.empty(length)) + torch.nn.init.constant_(biases, bias_start) + self._biases_dict[length] = biases + self._rnn_network.register_parameter('{}_biases_{}'.format(self._type, str(length)), + biases) + + return self._biases_dict[length] + + +class GMSDRCell(torch.nn.Module): + def __init__(self, num_units, input_dim, adj_mx, max_diffusion_step, num_nodes, pre_k, pre_v, nonlinearity='tanh', + use_gc_for_ru=True): + """ + + :param num_units: + :param adj_mx: + :param max_diffusion_step: + :param num_nodes: + :param nonlinearity: + :param filter_type: "laplacian", "random_walk", "dual_random_walk". + :param use_gc_for_ru: whether to use Graph convolution to calculate the reset and update gates. + """ + + super().__init__() + self._activation = torch.tanh if nonlinearity == 'tanh' else torch.relu + # support other nonlinearities up here? + self._num_nodes = num_nodes + self._num_units = num_units + self._max_diffusion_step = max_diffusion_step + self._supports = [] + self._use_gc_for_ru = use_gc_for_ru + self.pre_k = pre_k + self.pre_v = pre_v + self.input_dim = input_dim + self.nodevec1 = nn.Parameter(torch.randn(num_nodes, 10), requires_grad=True) + self.nodevec2 = nn.Parameter(torch.randn(10, num_nodes), requires_grad=True) + + # supports = [] + # if filter_type == "laplacian": + # supports.append(utils.calculate_scaled_laplacian(adj_mx, lambda_max=None)) + # elif filter_type == "random_walk": + # supports.append(utils.calculate_random_walk_matrix(adj_mx).T) + # elif filter_type == "dual_random_walk": + # supports.append(utils.calculate_random_walk_matrix(adj_mx).T) + # supports.append(utils.calculate_random_walk_matrix(adj_mx.T).T) + # else: + # supports.append(utils.calculate_scaled_laplacian(adj_mx)) + # for support in supports: + # self._supports.append(support) + # # self._supports.append(self._build_sparse_matrix(support)) + self._supports = adj_mx + + self._fc_params = LayerParams(self, 'fc') + self._gconv_params = LayerParams(self, 'gconv') + self.W = nn.Parameter(torch.zeros(self._num_units, self._num_units), requires_grad=True) + self.b = nn.Parameter(torch.zeros(num_nodes, self._num_units), requires_grad=True) + self.R = nn.Parameter(torch.zeros(pre_k, num_nodes, self._num_units), requires_grad=True) + self.attlinear = nn.Linear(num_nodes * self._num_units, 1) + + # @staticmethod + # def _build_sparse_matrix(L): + # L = L.tocoo() + # indices = np.column_stack((L.row, L.col)) + # # this is to ensure row-major ordering to equal torch.sparse.sparse_reorder(L) + # indices = indices[np.lexsort((indices[:, 0], indices[:, 1]))] + # L = torch.sparse_coo_tensor(indices.T, L.data, L.shape, device=device) + # return L + + def forward(self, inputs, hx_k): + """Gated recurrent unit (GRU) with Graph Convolution. + :param inputs: (B, num_nodes * input_dim) + :param hx_k: (B, pre_k, num_nodes, rnn_units) + + :return + - Output: A `2-D` tensor with shape `(B, num_nodes * rnn_units)`. + """ + bs, k, n, d = hx_k.shape + preH = hx_k[:, -1:] + for i in range(1, self.pre_v): + preH = torch.cat([preH, hx_k[:, -(i + 1):-i]], -1) + preH = preH.reshape(bs, n, d * self.pre_v) + self.adp = F.softmax(F.relu(torch.mm(self.nodevec1, self.nodevec2)), dim=1) + convInput = F.leaky_relu_(self._gconv(inputs, preH, d, bias_start=1.0)) + new_states = hx_k + self.R.unsqueeze(0) + output = torch.matmul(convInput, self.W) + self.b.unsqueeze(0) + self.attention(new_states) + output = output.unsqueeze(1) + x = hx_k[:, 1:k] + hx_k = torch.cat([x, output], dim=1) + output = output.reshape(bs, n * d) + return output, hx_k + + @staticmethod + def _concat(x, x_): + x_ = x_.unsqueeze(0) + return torch.cat([x, x_], dim=0) + + def _gconv(self, inputs, state, output_size, bias_start=0.0): + # input / state: (batch_size, num_nodes, input_dim/state_dim) + batch_size = inputs.shape[0] + inputs = torch.reshape(inputs, (batch_size, self._num_nodes, -1)) + state = torch.reshape(state, (batch_size, self._num_nodes, -1)) + inputs_and_state = torch.cat([inputs, state], dim=2) + input_size = inputs_and_state.size(2) + + x = inputs_and_state + x0 = x.permute(1, 2, 0) # (num_nodes, total_arg_size, batch_size) + x0 = torch.reshape(x0, shape=[self._num_nodes, input_size * batch_size]) + x = torch.unsqueeze(x0, 0) + if self._max_diffusion_step == 0: + pass + else: + for support in self._supports: + x1 = torch.mm(support.to(x0.device), x0) + x = self._concat(x, x1) + + for k in range(2, self._max_diffusion_step + 1): + x2 = 2 * torch.mm(support.to(x0.device), x1) - x0 + x = self._concat(x, x2) + x1, x0 = x2, x1 + x1 = self.adp.mm(x0) + x = self._concat(x, x1) + for k in range(2, self._max_diffusion_step + 1): + x2 = self.adp.mm(x1) - x0 + x = self._concat(x, x2) + x1, x0 = x2, x1 + num_matrices = (len(self._supports) + 1) * self._max_diffusion_step + 1 + # num_matrices = (len(self._supports)) * self._max_diffusion_step + 1 + x = torch.reshape(x, shape=[num_matrices, self._num_nodes, input_size, batch_size]) + x = x.permute(3, 1, 2, 0) # (batch_size, num_nodes, input_size, order) + x = torch.reshape(x, shape=[batch_size * self._num_nodes, input_size * num_matrices]) + + weights = self._gconv_params.get_weights((input_size * num_matrices, output_size)).to(x.device) + x = torch.matmul(x, weights) # (batch_size * self._num_nodes, output_size) + + biases = self._gconv_params.get_biases(output_size, bias_start).to(x.device) + x += biases + # Reshape res back to 2D: (batch_size, num_node, state_dim) -> (batch_size, num_node * state_dim) + return torch.reshape(x, [batch_size, self._num_nodes, output_size]) + + def attention(self, inputs: Tensor): + bs, k, n, d = inputs.size() + x = inputs.reshape(bs, k, -1) + out = self.attlinear(x) + weight = F.softmax(out, dim=1) + outputs = (x * weight).sum(dim=1).reshape(bs, n, d) + return outputs \ No newline at end of file diff --git a/baselines/GMSDR/run.sh b/baselines/GMSDR/run.sh new file mode 100644 index 00000000..2f46b383 --- /dev/null +++ b/baselines/GMSDR/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash +python experiments/train.py -c baselines/GMSDR/METR-LA.py --gpus '4' +python experiments/train.py -c baselines/GMSDR/PEMS-BAY.py --gpus '4' +python experiments/train.py -c baselines/GMSDR/PEMS03.py --gpus '4' +python experiments/train.py -c baselines/GMSDR/PEMS04.py --gpus '4' +python experiments/train.py -c baselines/GMSDR/PEMS07.py --gpus '4' +python experiments/train.py -c baselines/GMSDR/PEMS08.py --gpus '4' \ No newline at end of file