Skip to content

Commit 28e635b

Browse files
superarthurlxsuperarthurlx
committed
'add-s-mamba'
1 parent 962435e commit 28e635b

9 files changed

Lines changed: 856 additions & 94 deletions

File tree

‎baselines/S_Mamba/ETTm2.py‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
import os
2+
import sys
3+
from easydict import EasyDict
4+
sys.path.append(os.path.abspath(__file__ + '/../../..'))
5+
6+
from basicts.metrics import masked_mae, masked_mse
7+
from basicts.data import TimeSeriesForecastingDataset
8+
from basicts.runners import SimpleTimeSeriesForecastingRunner
9+
from basicts.scaler import ZScoreScaler
10+
from basicts.utils import get_regular_settings
11+
12+
from .arch import S_Mamba
13+
14+
############################## Hot Parameters ##############################
15+
# Dataset & Metrics configuration
16+
DATA_NAME = 'ETTm2' # Dataset name
17+
regular_settings = get_regular_settings(DATA_NAME)
18+
INPUT_LEN = regular_settings['INPUT_LEN'] # Length of input sequence
19+
OUTPUT_LEN = regular_settings['OUTPUT_LEN'] # Length of output sequence
20+
TRAIN_VAL_TEST_RATIO = regular_settings['TRAIN_VAL_TEST_RATIO'] # Train/Validation/Test split ratios
21+
NORM_EACH_CHANNEL = regular_settings['NORM_EACH_CHANNEL'] # Whether to normalize each channel of the data
22+
RESCALE = regular_settings['RESCALE'] # Whether to rescale the data
23+
NULL_VAL = regular_settings['NULL_VAL'] # Null value in the data
24+
# Model architecture and parameters
25+
MODEL_ARCH = S_Mamba
26+
NUM_NODES = 7
27+
MODEL_PARAM = {
28+
"enc_in": NUM_NODES, # num nodes
29+
"seq_len": INPUT_LEN,
30+
"pred_len": OUTPUT_LEN, # prediction sequence length
31+
"d_model": 256,
32+
"d_state": 2,
33+
"d_ff": 256,
34+
"e_layers": 2,
35+
"use_norm": True,
36+
"embed": "timeF", # [timeF, fixed, learned]
37+
"freq": 'h',
38+
"dropout": 0.1,
39+
"activation": "gelu",
40+
"num_time_features": 4, # number of used time features
41+
"time_of_day_size": 96,
42+
"day_of_week_size": 7,
43+
"day_of_month_size": 31,
44+
"day_of_year_size": 366
45+
}
46+
47+
NUM_EPOCHS = 50
48+
49+
############################## General Configuration ##############################
50+
CFG = EasyDict()
51+
# General settings
52+
CFG.DESCRIPTION = 'An Example Config'
53+
CFG.GPU_NUM = 1 # Number of GPUs to use (0 for CPU mode)
54+
# Runner
55+
CFG.RUNNER = SimpleTimeSeriesForecastingRunner
56+
57+
############################## Environment Configuration ##############################
58+
CFG.ENV = EasyDict() # Environment settings. Default: None
59+
CFG.ENV.SEED = 1 # Random seed. Default: None
60+
61+
############################## Dataset Configuration ##############################
62+
CFG.DATASET = EasyDict()
63+
# Dataset settings
64+
CFG.DATASET.NAME = DATA_NAME
65+
CFG.DATASET.TYPE = TimeSeriesForecastingDataset
66+
CFG.DATASET.PARAM = EasyDict({
67+
'dataset_name': DATA_NAME,
68+
'train_val_test_ratio': TRAIN_VAL_TEST_RATIO,
69+
'input_len': INPUT_LEN,
70+
'output_len': OUTPUT_LEN,
71+
# 'mode' is automatically set by the runner
72+
})
73+
74+
############################## Scaler Configuration ##############################
75+
CFG.SCALER = EasyDict()
76+
# Scaler settings
77+
CFG.SCALER.TYPE = ZScoreScaler # Scaler class
78+
CFG.SCALER.PARAM = EasyDict({
79+
'dataset_name': DATA_NAME,
80+
'train_ratio': TRAIN_VAL_TEST_RATIO[0],
81+
'norm_each_channel': NORM_EACH_CHANNEL,
82+
'rescale': RESCALE,
83+
})
84+
85+
############################## Model Configuration ##############################
86+
CFG.MODEL = EasyDict()
87+
# Model settings
88+
CFG.MODEL.NAME = MODEL_ARCH.__name__
89+
CFG.MODEL.ARCH = MODEL_ARCH
90+
CFG.MODEL.PARAM = MODEL_PARAM
91+
CFG.MODEL.FORWARD_FEATURES = [0, 1, 2, 3, 4]
92+
CFG.MODEL.TARGET_FEATURES = [0]
93+
94+
############################## Metrics Configuration ##############################
95+
96+
CFG.METRICS = EasyDict()
97+
# Metrics settings
98+
CFG.METRICS.FUNCS = EasyDict({
99+
'MAE': masked_mae,
100+
'MSE': masked_mse,
101+
})
102+
CFG.METRICS.TARGET = 'MSE'
103+
CFG.METRICS.NULL_VAL = NULL_VAL
104+
105+
############################## Training Configuration ##############################
106+
CFG.TRAIN = EasyDict()
107+
CFG.TRAIN.NUM_EPOCHS = NUM_EPOCHS
108+
CFG.TRAIN.CKPT_SAVE_DIR = os.path.join(
109+
'checkpoints',
110+
MODEL_ARCH.__name__,
111+
'_'.join([DATA_NAME, str(CFG.TRAIN.NUM_EPOCHS), str(INPUT_LEN), str(OUTPUT_LEN)])
112+
)
113+
CFG.TRAIN.LOSS = masked_mse
114+
# Optimizer settings
115+
CFG.TRAIN.OPTIM = EasyDict()
116+
CFG.TRAIN.OPTIM.TYPE = "Adam"
117+
CFG.TRAIN.OPTIM.PARAM = {
118+
"lr": 0.001
119+
}
120+
# Learning rate scheduler settings
121+
CFG.TRAIN.LR_SCHEDULER = EasyDict()
122+
CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
123+
CFG.TRAIN.LR_SCHEDULER.PARAM = {
124+
"milestones": [1, 25]
125+
}
126+
CFG.TRAIN.CLIP_GRAD_PARAM = {
127+
'max_norm': 5.0
128+
}
129+
# Train data loader settings
130+
CFG.TRAIN.DATA = EasyDict()
131+
CFG.TRAIN.DATA.BATCH_SIZE = 64
132+
CFG.TRAIN.DATA.SHUFFLE = True
133+
134+
############################## Validation Configuration ##############################
135+
CFG.VAL = EasyDict()
136+
CFG.VAL.INTERVAL = 1
137+
CFG.VAL.DATA = EasyDict()
138+
CFG.VAL.DATA.BATCH_SIZE = 64
139+
140+
############################## Test Configuration ##############################
141+
CFG.TEST = EasyDict()
142+
CFG.TEST.INTERVAL = 1
143+
CFG.TEST.DATA = EasyDict()
144+
CFG.TEST.DATA.BATCH_SIZE = 64
145+
146+
############################## Evaluation Configuration ##############################
147+
148+
CFG.EVAL = EasyDict()
149+
150+
# Evaluation parameters
151+
CFG.EVAL.HORIZONS = [12, 24, 48, 96]
152+
CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True

‎baselines/S_Mamba/Electricity.py‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
import os
2+
import sys
3+
from easydict import EasyDict
4+
sys.path.append(os.path.abspath(__file__ + '/../../..'))
5+
6+
from basicts.metrics import masked_mae, masked_mse
7+
from basicts.data import TimeSeriesForecastingDataset
8+
from basicts.runners import SimpleTimeSeriesForecastingRunner
9+
from basicts.scaler import ZScoreScaler
10+
from basicts.utils import get_regular_settings
11+
12+
from .arch import S_Mamba
13+
14+
############################## Hot Parameters ##############################
15+
# Dataset & Metrics configuration
16+
DATA_NAME = 'Electricity' # Dataset name
17+
regular_settings = get_regular_settings(DATA_NAME)
18+
INPUT_LEN = regular_settings['INPUT_LEN'] # Length of input sequence
19+
OUTPUT_LEN = regular_settings['OUTPUT_LEN'] # Length of output sequence
20+
TRAIN_VAL_TEST_RATIO = regular_settings['TRAIN_VAL_TEST_RATIO'] # Train/Validation/Test split ratios
21+
NORM_EACH_CHANNEL = regular_settings['NORM_EACH_CHANNEL'] # Whether to normalize each channel of the data
22+
RESCALE = regular_settings['RESCALE'] # Whether to rescale the data
23+
NULL_VAL = regular_settings['NULL_VAL'] # Null value in the data
24+
# Model architecture and parameters
25+
MODEL_ARCH = S_Mamba
26+
NUM_NODES = 321
27+
MODEL_PARAM = {
28+
"enc_in": NUM_NODES, # num nodes
29+
"seq_len": INPUT_LEN,
30+
"pred_len": OUTPUT_LEN, # prediction sequence length
31+
"d_model": 512,
32+
"d_state": 16,
33+
"d_ff": 512,
34+
"e_layers": 3,
35+
"use_norm": True,
36+
"embed": "timeF", # [timeF, fixed, learned]
37+
"freq": 'h',
38+
"dropout": 0.1,
39+
"activation": "gelu",
40+
"num_time_features": 4, # number of used time features
41+
"time_of_day_size": 24,
42+
"day_of_week_size": 7,
43+
"day_of_month_size": 31,
44+
"day_of_year_size": 366
45+
}
46+
47+
NUM_EPOCHS = 50
48+
49+
############################## General Configuration ##############################
50+
CFG = EasyDict()
51+
# General settings
52+
CFG.DESCRIPTION = 'An Example Config'
53+
CFG.GPU_NUM = 1 # Number of GPUs to use (0 for CPU mode)
54+
# Runner
55+
CFG.RUNNER = SimpleTimeSeriesForecastingRunner
56+
57+
############################## Environment Configuration ##############################
58+
CFG.ENV = EasyDict() # Environment settings. Default: None
59+
CFG.ENV.SEED = 1 # Random seed. Default: None
60+
61+
############################## Dataset Configuration ##############################
62+
CFG.DATASET = EasyDict()
63+
# Dataset settings
64+
CFG.DATASET.NAME = DATA_NAME
65+
CFG.DATASET.TYPE = TimeSeriesForecastingDataset
66+
CFG.DATASET.PARAM = EasyDict({
67+
'dataset_name': DATA_NAME,
68+
'train_val_test_ratio': TRAIN_VAL_TEST_RATIO,
69+
'input_len': INPUT_LEN,
70+
'output_len': OUTPUT_LEN,
71+
# 'mode' is automatically set by the runner
72+
})
73+
74+
############################## Scaler Configuration ##############################
75+
CFG.SCALER = EasyDict()
76+
# Scaler settings
77+
CFG.SCALER.TYPE = ZScoreScaler # Scaler class
78+
CFG.SCALER.PARAM = EasyDict({
79+
'dataset_name': DATA_NAME,
80+
'train_ratio': TRAIN_VAL_TEST_RATIO[0],
81+
'norm_each_channel': NORM_EACH_CHANNEL,
82+
'rescale': RESCALE,
83+
})
84+
85+
############################## Model Configuration ##############################
86+
CFG.MODEL = EasyDict()
87+
# Model settings
88+
CFG.MODEL.NAME = MODEL_ARCH.__name__
89+
CFG.MODEL.ARCH = MODEL_ARCH
90+
CFG.MODEL.PARAM = MODEL_PARAM
91+
CFG.MODEL.FORWARD_FEATURES = [0, 1, 2, 3, 4]
92+
CFG.MODEL.TARGET_FEATURES = [0]
93+
94+
############################## Metrics Configuration ##############################
95+
96+
CFG.METRICS = EasyDict()
97+
# Metrics settings
98+
CFG.METRICS.FUNCS = EasyDict({
99+
'MAE': masked_mae,
100+
'MSE': masked_mse,
101+
})
102+
CFG.METRICS.TARGET = 'MSE'
103+
CFG.METRICS.NULL_VAL = NULL_VAL
104+
105+
############################## Training Configuration ##############################
106+
CFG.TRAIN = EasyDict()
107+
CFG.TRAIN.NUM_EPOCHS = NUM_EPOCHS
108+
CFG.TRAIN.CKPT_SAVE_DIR = os.path.join(
109+
'checkpoints',
110+
MODEL_ARCH.__name__,
111+
'_'.join([DATA_NAME, str(CFG.TRAIN.NUM_EPOCHS), str(INPUT_LEN), str(OUTPUT_LEN)])
112+
)
113+
CFG.TRAIN.LOSS = masked_mse
114+
# Optimizer settings
115+
CFG.TRAIN.OPTIM = EasyDict()
116+
CFG.TRAIN.OPTIM.TYPE = "Adam"
117+
CFG.TRAIN.OPTIM.PARAM = {
118+
"lr": 0.001
119+
}
120+
# Learning rate scheduler settings
121+
CFG.TRAIN.LR_SCHEDULER = EasyDict()
122+
CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR"
123+
CFG.TRAIN.LR_SCHEDULER.PARAM = {
124+
"milestones": [1, 25]
125+
}
126+
CFG.TRAIN.CLIP_GRAD_PARAM = {
127+
'max_norm': 5.0
128+
}
129+
# Train data loader settings
130+
CFG.TRAIN.DATA = EasyDict()
131+
CFG.TRAIN.DATA.BATCH_SIZE = 64
132+
CFG.TRAIN.DATA.SHUFFLE = True
133+
134+
############################## Validation Configuration ##############################
135+
CFG.VAL = EasyDict()
136+
CFG.VAL.INTERVAL = 1
137+
CFG.VAL.DATA = EasyDict()
138+
CFG.VAL.DATA.BATCH_SIZE = 64
139+
140+
############################## Test Configuration ##############################
141+
CFG.TEST = EasyDict()
142+
CFG.TEST.INTERVAL = 1
143+
CFG.TEST.DATA = EasyDict()
144+
CFG.TEST.DATA.BATCH_SIZE = 64
145+
146+
############################## Evaluation Configuration ##############################
147+
148+
CFG.EVAL = EasyDict()
149+
150+
# Evaluation parameters
151+
CFG.EVAL.HORIZONS = [12, 24, 48, 96]
152+
CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True

0 commit comments

Comments
 (0)