From 28e635b9acfce38e671f6dc8fec0af39b69492dc Mon Sep 17 00:00:00 2001 From: lx Date: Sun, 15 Jun 2025 17:04:57 +0800 Subject: [PATCH] 'add-s-mamba' --- baselines/S_Mamba/ETTm2.py | 152 ++++++++++++++++++++ baselines/S_Mamba/Electricity.py | 152 ++++++++++++++++++++ baselines/S_Mamba/Weather.py | 152 ++++++++++++++++++++ baselines/S_Mamba/arch/Embed.py | 189 +++++++++++++++++++++++++ baselines/S_Mamba/arch/Mamba_EncDec.py | 69 +++++++++ baselines/S_Mamba/arch/__init__.py | 1 + baselines/S_Mamba/arch/masking.py | 26 ++++ baselines/S_Mamba/arch/s_mamba_arch.py | 115 +++++++++++++++ basicts/runners/optim/lr_schedulers.py | 94 ------------ 9 files changed, 856 insertions(+), 94 deletions(-) create mode 100644 baselines/S_Mamba/ETTm2.py create mode 100644 baselines/S_Mamba/Electricity.py create mode 100644 baselines/S_Mamba/Weather.py create mode 100644 baselines/S_Mamba/arch/Embed.py create mode 100644 baselines/S_Mamba/arch/Mamba_EncDec.py create mode 100644 baselines/S_Mamba/arch/__init__.py create mode 100644 baselines/S_Mamba/arch/masking.py create mode 100644 baselines/S_Mamba/arch/s_mamba_arch.py delete mode 100644 basicts/runners/optim/lr_schedulers.py diff --git a/baselines/S_Mamba/ETTm2.py b/baselines/S_Mamba/ETTm2.py new file mode 100644 index 00000000..22b8edb6 --- /dev/null +++ b/baselines/S_Mamba/ETTm2.py @@ -0,0 +1,152 @@ +import os +import sys +from easydict import EasyDict +sys.path.append(os.path.abspath(__file__ + '/../../..')) + +from basicts.metrics import masked_mae, masked_mse +from basicts.data import TimeSeriesForecastingDataset +from basicts.runners import SimpleTimeSeriesForecastingRunner +from basicts.scaler import ZScoreScaler +from basicts.utils import get_regular_settings + +from .arch import S_Mamba + +############################## Hot Parameters ############################## +# Dataset & Metrics configuration +DATA_NAME = 'ETTm2' # Dataset name +regular_settings = get_regular_settings(DATA_NAME) +INPUT_LEN = regular_settings['INPUT_LEN'] # Length of input sequence +OUTPUT_LEN = regular_settings['OUTPUT_LEN'] # Length of output sequence +TRAIN_VAL_TEST_RATIO = regular_settings['TRAIN_VAL_TEST_RATIO'] # Train/Validation/Test split ratios +NORM_EACH_CHANNEL = regular_settings['NORM_EACH_CHANNEL'] # Whether to normalize each channel of the data +RESCALE = regular_settings['RESCALE'] # Whether to rescale the data +NULL_VAL = regular_settings['NULL_VAL'] # Null value in the data +# Model architecture and parameters +MODEL_ARCH = S_Mamba +NUM_NODES = 7 +MODEL_PARAM = { + "enc_in": NUM_NODES, # num nodes + "seq_len": INPUT_LEN, + "pred_len": OUTPUT_LEN, # prediction sequence length + "d_model": 256, + "d_state": 2, + "d_ff": 256, + "e_layers": 2, + "use_norm": True, + "embed": "timeF", # [timeF, fixed, learned] + "freq": 'h', + "dropout": 0.1, + "activation": "gelu", + "num_time_features": 4, # number of used time features + "time_of_day_size": 96, + "day_of_week_size": 7, + "day_of_month_size": 31, + "day_of_year_size": 366 + } + +NUM_EPOCHS = 50 + +############################## General Configuration ############################## +CFG = EasyDict() +# General settings +CFG.DESCRIPTION = 'An Example Config' +CFG.GPU_NUM = 1 # Number of GPUs to use (0 for CPU mode) +# Runner +CFG.RUNNER = SimpleTimeSeriesForecastingRunner + +############################## Environment Configuration ############################## +CFG.ENV = EasyDict() # Environment settings. Default: None +CFG.ENV.SEED = 1 # Random seed. Default: None + +############################## Dataset Configuration ############################## +CFG.DATASET = EasyDict() +# Dataset settings +CFG.DATASET.NAME = DATA_NAME +CFG.DATASET.TYPE = TimeSeriesForecastingDataset +CFG.DATASET.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_val_test_ratio': TRAIN_VAL_TEST_RATIO, + 'input_len': INPUT_LEN, + 'output_len': OUTPUT_LEN, + # 'mode' is automatically set by the runner +}) + +############################## Scaler Configuration ############################## +CFG.SCALER = EasyDict() +# Scaler settings +CFG.SCALER.TYPE = ZScoreScaler # Scaler class +CFG.SCALER.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_ratio': TRAIN_VAL_TEST_RATIO[0], + 'norm_each_channel': NORM_EACH_CHANNEL, + 'rescale': RESCALE, +}) + +############################## Model Configuration ############################## +CFG.MODEL = EasyDict() +# Model settings +CFG.MODEL.NAME = MODEL_ARCH.__name__ +CFG.MODEL.ARCH = MODEL_ARCH +CFG.MODEL.PARAM = MODEL_PARAM +CFG.MODEL.FORWARD_FEATURES = [0, 1, 2, 3, 4] +CFG.MODEL.TARGET_FEATURES = [0] + +############################## Metrics Configuration ############################## + +CFG.METRICS = EasyDict() +# Metrics settings +CFG.METRICS.FUNCS = EasyDict({ + 'MAE': masked_mae, + 'MSE': masked_mse, + }) +CFG.METRICS.TARGET = 'MSE' +CFG.METRICS.NULL_VAL = NULL_VAL + +############################## Training Configuration ############################## +CFG.TRAIN = EasyDict() +CFG.TRAIN.NUM_EPOCHS = NUM_EPOCHS +CFG.TRAIN.CKPT_SAVE_DIR = os.path.join( + 'checkpoints', + MODEL_ARCH.__name__, + '_'.join([DATA_NAME, str(CFG.TRAIN.NUM_EPOCHS), str(INPUT_LEN), str(OUTPUT_LEN)]) +) +CFG.TRAIN.LOSS = masked_mse +# Optimizer settings +CFG.TRAIN.OPTIM = EasyDict() +CFG.TRAIN.OPTIM.TYPE = "Adam" +CFG.TRAIN.OPTIM.PARAM = { + "lr": 0.001 +} +# Learning rate scheduler settings +CFG.TRAIN.LR_SCHEDULER = EasyDict() +CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR" +CFG.TRAIN.LR_SCHEDULER.PARAM = { + "milestones": [1, 25] +} +CFG.TRAIN.CLIP_GRAD_PARAM = { + 'max_norm': 5.0 +} +# Train data loader settings +CFG.TRAIN.DATA = EasyDict() +CFG.TRAIN.DATA.BATCH_SIZE = 64 +CFG.TRAIN.DATA.SHUFFLE = True + +############################## Validation Configuration ############################## +CFG.VAL = EasyDict() +CFG.VAL.INTERVAL = 1 +CFG.VAL.DATA = EasyDict() +CFG.VAL.DATA.BATCH_SIZE = 64 + +############################## Test Configuration ############################## +CFG.TEST = EasyDict() +CFG.TEST.INTERVAL = 1 +CFG.TEST.DATA = EasyDict() +CFG.TEST.DATA.BATCH_SIZE = 64 + +############################## Evaluation Configuration ############################## + +CFG.EVAL = EasyDict() + +# Evaluation parameters +CFG.EVAL.HORIZONS = [12, 24, 48, 96] +CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True diff --git a/baselines/S_Mamba/Electricity.py b/baselines/S_Mamba/Electricity.py new file mode 100644 index 00000000..5924066f --- /dev/null +++ b/baselines/S_Mamba/Electricity.py @@ -0,0 +1,152 @@ +import os +import sys +from easydict import EasyDict +sys.path.append(os.path.abspath(__file__ + '/../../..')) + +from basicts.metrics import masked_mae, masked_mse +from basicts.data import TimeSeriesForecastingDataset +from basicts.runners import SimpleTimeSeriesForecastingRunner +from basicts.scaler import ZScoreScaler +from basicts.utils import get_regular_settings + +from .arch import S_Mamba + +############################## Hot Parameters ############################## +# Dataset & Metrics configuration +DATA_NAME = 'Electricity' # Dataset name +regular_settings = get_regular_settings(DATA_NAME) +INPUT_LEN = regular_settings['INPUT_LEN'] # Length of input sequence +OUTPUT_LEN = regular_settings['OUTPUT_LEN'] # Length of output sequence +TRAIN_VAL_TEST_RATIO = regular_settings['TRAIN_VAL_TEST_RATIO'] # Train/Validation/Test split ratios +NORM_EACH_CHANNEL = regular_settings['NORM_EACH_CHANNEL'] # Whether to normalize each channel of the data +RESCALE = regular_settings['RESCALE'] # Whether to rescale the data +NULL_VAL = regular_settings['NULL_VAL'] # Null value in the data +# Model architecture and parameters +MODEL_ARCH = S_Mamba +NUM_NODES = 321 +MODEL_PARAM = { + "enc_in": NUM_NODES, # num nodes + "seq_len": INPUT_LEN, + "pred_len": OUTPUT_LEN, # prediction sequence length + "d_model": 512, + "d_state": 16, + "d_ff": 512, + "e_layers": 3, + "use_norm": True, + "embed": "timeF", # [timeF, fixed, learned] + "freq": 'h', + "dropout": 0.1, + "activation": "gelu", + "num_time_features": 4, # number of used time features + "time_of_day_size": 24, + "day_of_week_size": 7, + "day_of_month_size": 31, + "day_of_year_size": 366 + } + +NUM_EPOCHS = 50 + +############################## General Configuration ############################## +CFG = EasyDict() +# General settings +CFG.DESCRIPTION = 'An Example Config' +CFG.GPU_NUM = 1 # Number of GPUs to use (0 for CPU mode) +# Runner +CFG.RUNNER = SimpleTimeSeriesForecastingRunner + +############################## Environment Configuration ############################## +CFG.ENV = EasyDict() # Environment settings. Default: None +CFG.ENV.SEED = 1 # Random seed. Default: None + +############################## Dataset Configuration ############################## +CFG.DATASET = EasyDict() +# Dataset settings +CFG.DATASET.NAME = DATA_NAME +CFG.DATASET.TYPE = TimeSeriesForecastingDataset +CFG.DATASET.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_val_test_ratio': TRAIN_VAL_TEST_RATIO, + 'input_len': INPUT_LEN, + 'output_len': OUTPUT_LEN, + # 'mode' is automatically set by the runner +}) + +############################## Scaler Configuration ############################## +CFG.SCALER = EasyDict() +# Scaler settings +CFG.SCALER.TYPE = ZScoreScaler # Scaler class +CFG.SCALER.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_ratio': TRAIN_VAL_TEST_RATIO[0], + 'norm_each_channel': NORM_EACH_CHANNEL, + 'rescale': RESCALE, +}) + +############################## Model Configuration ############################## +CFG.MODEL = EasyDict() +# Model settings +CFG.MODEL.NAME = MODEL_ARCH.__name__ +CFG.MODEL.ARCH = MODEL_ARCH +CFG.MODEL.PARAM = MODEL_PARAM +CFG.MODEL.FORWARD_FEATURES = [0, 1, 2, 3, 4] +CFG.MODEL.TARGET_FEATURES = [0] + +############################## Metrics Configuration ############################## + +CFG.METRICS = EasyDict() +# Metrics settings +CFG.METRICS.FUNCS = EasyDict({ + 'MAE': masked_mae, + 'MSE': masked_mse, + }) +CFG.METRICS.TARGET = 'MSE' +CFG.METRICS.NULL_VAL = NULL_VAL + +############################## Training Configuration ############################## +CFG.TRAIN = EasyDict() +CFG.TRAIN.NUM_EPOCHS = NUM_EPOCHS +CFG.TRAIN.CKPT_SAVE_DIR = os.path.join( + 'checkpoints', + MODEL_ARCH.__name__, + '_'.join([DATA_NAME, str(CFG.TRAIN.NUM_EPOCHS), str(INPUT_LEN), str(OUTPUT_LEN)]) +) +CFG.TRAIN.LOSS = masked_mse +# Optimizer settings +CFG.TRAIN.OPTIM = EasyDict() +CFG.TRAIN.OPTIM.TYPE = "Adam" +CFG.TRAIN.OPTIM.PARAM = { + "lr": 0.001 +} +# Learning rate scheduler settings +CFG.TRAIN.LR_SCHEDULER = EasyDict() +CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR" +CFG.TRAIN.LR_SCHEDULER.PARAM = { + "milestones": [1, 25] +} +CFG.TRAIN.CLIP_GRAD_PARAM = { + 'max_norm': 5.0 +} +# Train data loader settings +CFG.TRAIN.DATA = EasyDict() +CFG.TRAIN.DATA.BATCH_SIZE = 64 +CFG.TRAIN.DATA.SHUFFLE = True + +############################## Validation Configuration ############################## +CFG.VAL = EasyDict() +CFG.VAL.INTERVAL = 1 +CFG.VAL.DATA = EasyDict() +CFG.VAL.DATA.BATCH_SIZE = 64 + +############################## Test Configuration ############################## +CFG.TEST = EasyDict() +CFG.TEST.INTERVAL = 1 +CFG.TEST.DATA = EasyDict() +CFG.TEST.DATA.BATCH_SIZE = 64 + +############################## Evaluation Configuration ############################## + +CFG.EVAL = EasyDict() + +# Evaluation parameters +CFG.EVAL.HORIZONS = [12, 24, 48, 96] +CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True diff --git a/baselines/S_Mamba/Weather.py b/baselines/S_Mamba/Weather.py new file mode 100644 index 00000000..548bc0f3 --- /dev/null +++ b/baselines/S_Mamba/Weather.py @@ -0,0 +1,152 @@ +import os +import sys +from easydict import EasyDict +sys.path.append(os.path.abspath(__file__ + '/../../..')) + +from basicts.metrics import masked_mae, masked_mse +from basicts.data import TimeSeriesForecastingDataset +from basicts.runners import SimpleTimeSeriesForecastingRunner +from basicts.scaler import ZScoreScaler +from basicts.utils import get_regular_settings + +from .arch import S_Mamba + +############################## Hot Parameters ############################## +# Dataset & Metrics configuration +DATA_NAME = 'Weather' # Dataset name +regular_settings = get_regular_settings(DATA_NAME) +INPUT_LEN = regular_settings['INPUT_LEN'] # Length of input sequence +OUTPUT_LEN = regular_settings['OUTPUT_LEN'] # Length of output sequence +TRAIN_VAL_TEST_RATIO = regular_settings['TRAIN_VAL_TEST_RATIO'] # Train/Validation/Test split ratios +NORM_EACH_CHANNEL = regular_settings['NORM_EACH_CHANNEL'] # Whether to normalize each channel of the data +RESCALE = regular_settings['RESCALE'] # Whether to rescale the data +NULL_VAL = regular_settings['NULL_VAL'] # Null value in the data +# Model architecture and parameters +MODEL_ARCH = S_Mamba +NUM_NODES = 21 +MODEL_PARAM = { + "enc_in": NUM_NODES, # num nodes + "seq_len": INPUT_LEN, + "pred_len": OUTPUT_LEN, # prediction sequence length + "d_model": 512, + "d_state": 2, + "d_ff": 512, + "e_layers": 3, + "use_norm": True, + "embed": "timeF", # [timeF, fixed, learned] + "freq": 'h', + "dropout": 0.1, + "activation": "gelu", + "num_time_features": 4, # number of used time features + "time_of_day_size": 144, + "day_of_week_size": 7, + "day_of_month_size": 31, + "day_of_year_size": 366 + } + +NUM_EPOCHS = 50 + +############################## General Configuration ############################## +CFG = EasyDict() +# General settings +CFG.DESCRIPTION = 'An Example Config' +CFG.GPU_NUM = 1 # Number of GPUs to use (0 for CPU mode) +# Runner +CFG.RUNNER = SimpleTimeSeriesForecastingRunner + +############################## Environment Configuration ############################## +CFG.ENV = EasyDict() # Environment settings. Default: None +CFG.ENV.SEED = 1 # Random seed. Default: None + +############################## Dataset Configuration ############################## +CFG.DATASET = EasyDict() +# Dataset settings +CFG.DATASET.NAME = DATA_NAME +CFG.DATASET.TYPE = TimeSeriesForecastingDataset +CFG.DATASET.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_val_test_ratio': TRAIN_VAL_TEST_RATIO, + 'input_len': INPUT_LEN, + 'output_len': OUTPUT_LEN, + # 'mode' is automatically set by the runner +}) + +############################## Scaler Configuration ############################## +CFG.SCALER = EasyDict() +# Scaler settings +CFG.SCALER.TYPE = ZScoreScaler # Scaler class +CFG.SCALER.PARAM = EasyDict({ + 'dataset_name': DATA_NAME, + 'train_ratio': TRAIN_VAL_TEST_RATIO[0], + 'norm_each_channel': NORM_EACH_CHANNEL, + 'rescale': RESCALE, +}) + +############################## Model Configuration ############################## +CFG.MODEL = EasyDict() +# Model settings +CFG.MODEL.NAME = MODEL_ARCH.__name__ +CFG.MODEL.ARCH = MODEL_ARCH +CFG.MODEL.PARAM = MODEL_PARAM +CFG.MODEL.FORWARD_FEATURES = [0, 1, 2, 3, 4] +CFG.MODEL.TARGET_FEATURES = [0] + +############################## Metrics Configuration ############################## + +CFG.METRICS = EasyDict() +# Metrics settings +CFG.METRICS.FUNCS = EasyDict({ + 'MAE': masked_mae, + 'MSE': masked_mse, + }) +CFG.METRICS.TARGET = 'MSE' +CFG.METRICS.NULL_VAL = NULL_VAL + +############################## Training Configuration ############################## +CFG.TRAIN = EasyDict() +CFG.TRAIN.NUM_EPOCHS = NUM_EPOCHS +CFG.TRAIN.CKPT_SAVE_DIR = os.path.join( + 'checkpoints', + MODEL_ARCH.__name__, + '_'.join([DATA_NAME, str(CFG.TRAIN.NUM_EPOCHS), str(INPUT_LEN), str(OUTPUT_LEN)]) +) +CFG.TRAIN.LOSS = masked_mse +# Optimizer settings +CFG.TRAIN.OPTIM = EasyDict() +CFG.TRAIN.OPTIM.TYPE = "Adam" +CFG.TRAIN.OPTIM.PARAM = { + "lr": 0.001 +} +# Learning rate scheduler settings +CFG.TRAIN.LR_SCHEDULER = EasyDict() +CFG.TRAIN.LR_SCHEDULER.TYPE = "MultiStepLR" +CFG.TRAIN.LR_SCHEDULER.PARAM = { + "milestones": [1, 25] +} +CFG.TRAIN.CLIP_GRAD_PARAM = { + 'max_norm': 5.0 +} +# Train data loader settings +CFG.TRAIN.DATA = EasyDict() +CFG.TRAIN.DATA.BATCH_SIZE = 64 +CFG.TRAIN.DATA.SHUFFLE = True + +############################## Validation Configuration ############################## +CFG.VAL = EasyDict() +CFG.VAL.INTERVAL = 1 +CFG.VAL.DATA = EasyDict() +CFG.VAL.DATA.BATCH_SIZE = 64 + +############################## Test Configuration ############################## +CFG.TEST = EasyDict() +CFG.TEST.INTERVAL = 1 +CFG.TEST.DATA = EasyDict() +CFG.TEST.DATA.BATCH_SIZE = 64 + +############################## Evaluation Configuration ############################## + +CFG.EVAL = EasyDict() + +# Evaluation parameters +CFG.EVAL.HORIZONS = [12, 24, 48, 96] +CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True diff --git a/baselines/S_Mamba/arch/Embed.py b/baselines/S_Mamba/arch/Embed.py new file mode 100644 index 00000000..6bc2e72a --- /dev/null +++ b/baselines/S_Mamba/arch/Embed.py @@ -0,0 +1,189 @@ +import torch +import torch.nn as nn +import math + +class PositionalEmbedding(nn.Module): + def __init__(self, d_model, max_len=5000): + super(PositionalEmbedding, self).__init__() + # Compute the positional encodings once in log space. + pe = torch.zeros(max_len, d_model).float() + pe.require_grad = False + + position = torch.arange(0, max_len).float().unsqueeze(1) + div_term = (torch.arange(0, d_model, 2).float() + * -(math.log(10000.0) / d_model)).exp() + + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + + pe = pe.unsqueeze(0) + self.register_buffer('pe', pe) + + def forward(self, x): + return self.pe[:, :x.size(1)] + + +class TokenEmbedding(nn.Module): + def __init__(self, c_in, d_model): + super(TokenEmbedding, self).__init__() + padding = 1 if torch.__version__ >= '1.5.0' else 2 + self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model, + kernel_size=3, padding=padding, padding_mode='circular', bias=False) + for m in self.modules(): + if isinstance(m, nn.Conv1d): + nn.init.kaiming_normal_( + m.weight, mode='fan_in', nonlinearity='leaky_relu') + + def forward(self, x): + x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2) + return x + + +class FixedEmbedding(nn.Module): + def __init__(self, c_in, d_model): + super(FixedEmbedding, self).__init__() + + w = torch.zeros(c_in, d_model).float() + w.require_grad = False + + position = torch.arange(0, c_in).float().unsqueeze(1) + div_term = (torch.arange(0, d_model, 2).float() + * -(math.log(10000.0) / d_model)).exp() + + w[:, 0::2] = torch.sin(position * div_term) + w[:, 1::2] = torch.cos(position * div_term) + + self.emb = nn.Embedding(c_in, d_model) + self.emb.weight = nn.Parameter(w, requires_grad=False) + + def forward(self, x): + return self.emb(x).detach() + + +class TemporalEmbedding(nn.Module): + def __init__(self, d_model, embed_type='fixed', freq='h'): + super(TemporalEmbedding, self).__init__() + + minute_size = 4 + hour_size = 24 + weekday_size = 7 + day_size = 32 + month_size = 13 + + Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding + if freq == 't': + self.minute_embed = Embed(minute_size, d_model) + self.hour_embed = Embed(hour_size, d_model) + self.weekday_embed = Embed(weekday_size, d_model) + self.day_embed = Embed(day_size, d_model) + self.month_embed = Embed(month_size, d_model) + + def forward(self, x): + x = x.long() + minute_x = self.minute_embed(x[:, :, 4]) if hasattr( + self, 'minute_embed') else 0. + hour_x = self.hour_embed(x[:, :, 3]) + weekday_x = self.weekday_embed(x[:, :, 2]) + day_x = self.day_embed(x[:, :, 1]) + month_x = self.month_embed(x[:, :, 0]) + + return hour_x + weekday_x + day_x + month_x + minute_x + + +class TimeFeatureEmbedding(nn.Module): + def __init__(self, d_model, embed_type='timeF', freq='h'): + super(TimeFeatureEmbedding, self).__init__() + + freq_map = {'h': 4, 't': 5, 's': 6, + 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3} + d_inp = freq_map[freq] + self.embed = nn.Linear(d_inp, d_model, bias=False) + + def forward(self, x): + return self.embed(x) + + +class DataEmbedding(nn.Module): + def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1): + super(DataEmbedding, self).__init__() + + self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model) + self.position_embedding = PositionalEmbedding(d_model=d_model) + self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type, + freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding( + d_model=d_model, embed_type=embed_type, freq=freq) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + if x_mark is None: + x = self.value_embedding(x) + self.position_embedding(x) + else: + x = self.value_embedding( + x) + self.temporal_embedding(x_mark) + self.position_embedding(x) + return self.dropout(x) + +class moving_avg(nn.Module): + """ + Moving average block to highlight the trend of time series + """ + def __init__(self, kernel_size, stride): + super(moving_avg, self).__init__() + self.kernel_size = kernel_size + self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0) + + def forward(self, x): + # padding on the both ends of time series + front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1) + end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1) + x = torch.cat([front, x, end], dim=1) + x = self.avg(x.permute(0, 2, 1)) + x = x.permute(0, 2, 1) + return x + + +class series_decomp(nn.Module): + """ + Series decomposition block + """ + def __init__(self, kernel_size): + super(series_decomp, self).__init__() + self.moving_avg = moving_avg(kernel_size, stride=1) + + def forward(self, x): + moving_mean = self.moving_avg(x) + res = x - moving_mean + return res, moving_mean +class DataEmbedding_wo_pos(nn.Module): + def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1): + super(DataEmbedding_wo_pos, self).__init__() + + self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model) + self.position_embedding = PositionalEmbedding(d_model=d_model) + self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type, + freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding( + d_model=d_model, embed_type=embed_type, freq=freq) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + if x_mark == None: + x = self.value_embedding(x) + else: + x = self.value_embedding(x) + self.temporal_embedding(x_mark) + return self.dropout(x) +class DataEmbedding_inverted(nn.Module): + def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1): + super(DataEmbedding_inverted, self).__init__() + self.value_embedding = nn.Linear(c_in, d_model) + self.dropout = nn.Dropout(p=dropout) + + def forward(self, x, x_mark): + x = x.permute(0, 2, 1) + + # x: [Batch Variate Time] + if x_mark is None: + + x = self.value_embedding(x) + else: + x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], 1)) + return self.dropout(x) + diff --git a/baselines/S_Mamba/arch/Mamba_EncDec.py b/baselines/S_Mamba/arch/Mamba_EncDec.py new file mode 100644 index 00000000..75a7b3f1 --- /dev/null +++ b/baselines/S_Mamba/arch/Mamba_EncDec.py @@ -0,0 +1,69 @@ +import torch.nn as nn +import torch.nn.functional as F +from mamba_ssm import Mamba + +class EncoderLayer(nn.Module): + def __init__(self, attention, attention_r, d_model, d_ff=None, dropout=0.1, activation="relu"): + super(EncoderLayer, self).__init__() + d_ff = d_ff or 4 * d_model + self.attention = attention + self.attention_r = attention_r + self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1) + self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1) + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout = nn.Dropout(dropout) + self.activation = F.relu if activation == "relu" else F.gelu + self.man = Mamba( + d_model=11, # Model dimension d_model + d_state=16, # SSM state expansion factor + d_conv=2, # Local convolution width + expand=1, # Block expansion factor) + ) + self.man2 = Mamba( + d_model=11, # Model dimension d_model + d_state=16, # SSM state expansion factor + d_conv=2, # Local convolution width + expand=1, # Block expansion factor) + ) + + def forward(self, x, attn_mask=None, tau=None, delta=None): + new_x = self.attention(x) + self.attention_r(x.flip(dims=[1])).flip(dims=[1]) + attn = 1 + + x = x + new_x + y = x = self.norm1(x) + y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1)))) + y = self.dropout(self.conv2(y).transpose(-1, 1)) + + return self.norm2(x + y), attn + + +class Encoder(nn.Module): + def __init__(self, attn_layers, conv_layers=None, norm_layer=None): + super(Encoder, self).__init__() + self.attn_layers = nn.ModuleList(attn_layers) + self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None + self.norm = norm_layer + + def forward(self, x, attn_mask=None, tau=None, delta=None): + # x [B, L, D] + attns = [] + if self.conv_layers is not None: + for i, (attn_layer, conv_layer) in enumerate(zip(self.attn_layers, self.conv_layers)): + delta = delta if i == 0 else None + x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta) + x = conv_layer(x) + attns.append(attn) + x, attn = self.attn_layers[-1](x, tau=tau, delta=None) + attns.append(attn) + else: + for attn_layer in self.attn_layers: + x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta) + attns.append(attn) + + if self.norm is not None: + x = self.norm(x) + + return x, attns + diff --git a/baselines/S_Mamba/arch/__init__.py b/baselines/S_Mamba/arch/__init__.py new file mode 100644 index 00000000..3d432eea --- /dev/null +++ b/baselines/S_Mamba/arch/__init__.py @@ -0,0 +1 @@ +from .s_mamba_arch import S_Mamba diff --git a/baselines/S_Mamba/arch/masking.py b/baselines/S_Mamba/arch/masking.py new file mode 100644 index 00000000..a19cbf63 --- /dev/null +++ b/baselines/S_Mamba/arch/masking.py @@ -0,0 +1,26 @@ +import torch + + +class TriangularCausalMask(): + def __init__(self, B, L, device="cpu"): + mask_shape = [B, 1, L, L] + with torch.no_grad(): + self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device) + + @property + def mask(self): + return self._mask + + +class ProbMask(): + def __init__(self, B, H, L, index, scores, device="cpu"): + _mask = torch.ones(L, scores.shape[-1], dtype=torch.bool).to(device).triu(1) + _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1]) + indicator = _mask_ex[torch.arange(B)[:, None, None], + torch.arange(H)[None, :, None], + index, :].to(device) + self._mask = indicator.view(scores.shape).to(device) + + @property + def mask(self): + return self._mask diff --git a/baselines/S_Mamba/arch/s_mamba_arch.py b/baselines/S_Mamba/arch/s_mamba_arch.py new file mode 100644 index 00000000..ae3bf231 --- /dev/null +++ b/baselines/S_Mamba/arch/s_mamba_arch.py @@ -0,0 +1,115 @@ +import torch +import torch.nn as nn +from .Mamba_EncDec import Encoder, EncoderLayer +from .Embed import DataEmbedding_inverted +from mamba_ssm import Mamba +from basicts.utils import data_transformation_4_xformer + +import pdb + + +class S_Mamba(nn.Module): + """ + Paper: Is Mamba Effective for Time Series Forecasting? + Official Code: https://github.com/wzhwzhwzh0921/S-D-Mamba + Link: https://arxiv.org/abs/2403.11144v3 + Venue: Neurocomputing + Task: Long-term Time Series Forecasting + """ + def __init__(self, **model_args): + super(S_Mamba, self).__init__() + + self.embed = model_args['embed'] + self.freq = model_args["freq"] + self.pred_len = model_args['pred_len'] + self.seq_len = model_args['seq_len'] + self.use_norm =model_args['use_norm'] + self.d_ff = model_args['d_ff'] + self.d_model = model_args['d_model'] + self.d_state = model_args['d_state'] + self.dropout = model_args["dropout"] + self.activation = model_args['activation'] + self.e_layers = model_args['e_layers'] + + + # Embedding + self.enc_embedding = DataEmbedding_inverted(self.seq_len, self.d_model, self.embed, self.freq, + self.dropout) + + # Encoder-only architecture + self.encoder = Encoder( + [ + EncoderLayer( + Mamba( + d_model=self.d_model, # Model dimension d_model + d_state=self.d_state, # SSM state expansion factor + d_conv=2, # Local convolution width + expand=1, # Block expansion factor) + ), + Mamba( + d_model=self.d_model, # Model dimension d_model + d_state=self.d_state, # SSM state expansion factor + d_conv=2, # Local convolution width + expand=1, # Block expansion factor) + ), + self.d_model, + self.d_ff, + dropout=self.dropout, + activation=self.activation + ) for l in range(self.e_layers) + ], + norm_layer=torch.nn.LayerNorm(self.d_model) + ) + self.projector = nn.Linear(self.d_model, self.pred_len, bias=True) + + self.time_of_day_size = model_args['time_of_day_size'] + + def forward_xformer(self, x_enc: torch.Tensor, x_mark_enc: torch.Tensor, x_dec: torch.Tensor, x_mark_dec: torch.Tensor) -> torch.Tensor: + + if self.use_norm: + # Normalization from Non-stationary Transformer + means = x_enc.mean(1, keepdim=True).detach() + x_enc = x_enc - means + stdev = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5) + x_enc /= stdev + + _, _, N = x_enc.shape # B L N + # B: batch_size; E: d_model; + # L: seq_len; S: pred_len; + # N: number of variate (tokens), can also includes covariates + + # Embedding + # B L N -> B N E (B L N -> B L E in the vanilla Transformer) + enc_out = self.enc_embedding(x_enc, x_mark_enc) # covariates (e.g timestamp) can be also embedded as tokens + + # B N E -> B N E (B L E -> B L E in the vanilla Transformer) + # the dimensions of embedded time series has been inverted, and then processed by native attn, layernorm and ffn modules + enc_out, attns = self.encoder(enc_out, attn_mask=None) + # B N E -> B N S -> B S N + dec_out = self.projector(enc_out).permute(0, 2, 1)[:, :, :N] # filter the covariates + + if self.use_norm: + # De-Normalization from Non-stationary Transformer + dec_out = dec_out * (stdev[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1)) + dec_out = dec_out + (means[:, 0, :].unsqueeze(1).repeat(1, self.pred_len, 1)) + + return dec_out + + def forward(self, history_data: torch.Tensor, future_data: torch.Tensor, batch_seen: int, epoch: int, train: bool, + **kwargs) -> torch.Tensor: + """ + + Args: + history_data (Tensor): Input data with shape: [B, L1, N, C] + future_data (Tensor): Future data with shape: [B, L2, N, C] + + Returns: + torch.Tensor: outputs with shape [B, L2, N, 1] + """ + history_data[..., 1] = history_data[..., 1] * self.time_of_day_size // (self.time_of_day_size / 24) / 23.0 + x_enc, x_mark_enc, x_dec, x_mark_dec = data_transformation_4_xformer(history_data=history_data, + future_data=future_data, + start_token_len=0) + #print(x_mark_enc.shape, x_mark_dec.shape) + prediction = self.forward_xformer(x_enc=x_enc, x_mark_enc=x_mark_enc, x_dec=x_dec, x_mark_dec=x_mark_dec) + return prediction.unsqueeze(-1) \ No newline at end of file diff --git a/basicts/runners/optim/lr_schedulers.py b/basicts/runners/optim/lr_schedulers.py deleted file mode 100644 index 6bafe92c..00000000 --- a/basicts/runners/optim/lr_schedulers.py +++ /dev/null @@ -1,94 +0,0 @@ -# define more learning rate shedulers here - -import math -from functools import partial - -from torch.optim import Optimizer -from torch.optim.lr_scheduler import LambdaLR - -__all__ = ['CosineWarmup', 'CosineWarmupRestarts'] - - -class CosineWarmup(LambdaLR): - """ - Create a schedule with a learning rate that decreases following the values of the cosine function between the - initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the - initial lr set in the optimizer. - - Modified from https://github.com/huggingface/transformers/blob/v4.46.0/src/transformers/optimization.py#L144 - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - num_cycles (`float`, *optional*, defaults to 0.5): - The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 - following a half-cosine). - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - def __init__(self, optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1): - lr_lambda = partial( - self._get_cosine_schedule_with_warmup_lr_lambda, - num_warmup_steps=num_warmup_steps, - num_training_steps=num_training_steps, - num_cycles=num_cycles, - ) - super().__init__(optimizer, lr_lambda, last_epoch) - - @staticmethod - def _get_cosine_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) - return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) - - -class CosineWarmupRestarts(LambdaLR): - """ - Create a schedule with a learning rate that decreases following the values of the cosine function between the - initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases - linearly between 0 and the initial lr set in the optimizer. - - # Modified from https://github.com/huggingface/transformers/blob/c2820c94916e34baf4486accae74760972183a2f/src/transformers/optimization.py#L144 - - Args: - optimizer ([`~torch.optim.Optimizer`]): - The optimizer for which to schedule the learning rate. - num_warmup_steps (`int`): - The number of steps for the warmup phase. - num_training_steps (`int`): - The total number of training steps. - num_cycles (`int`, *optional*, defaults to 1): - The number of hard restarts to use. - last_epoch (`int`, *optional*, defaults to -1): - The index of the last epoch when resuming training. - - Return: - `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. - """ - def __init__(self, optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1): - lr_lambda = partial( - self._get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda, - num_warmup_steps=num_warmup_steps, - num_training_steps=num_training_steps, - num_cycles=num_cycles, - ) - super().__init__(optimizer, lr_lambda, last_epoch) - - @staticmethod - def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda( - current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int - ): - if current_step < num_warmup_steps: - return float(current_step) / float(max(1, num_warmup_steps)) - progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) - if progress >= 1.0: - return 0.0 - return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))