added code
This commit is contained in:
@@ -0,0 +1,279 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.config_utils import *
|
||||
from configs.feature_config import feature_config
|
||||
from experiment_setup import get_fertility_based_eval_functions
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.cnn_lstm import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
cnn_lstm_base_config = {
|
||||
"model_name": "cnn_lstm",
|
||||
"version": "1.0.0",
|
||||
"model_class": CNNLSTM,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
}
|
||||
|
||||
base_cnn_lstm_training_config = {
|
||||
"batch_size": 128,
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-5,
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="cnn_lstm_complexity_run",
|
||||
base_model_configuration=cnn_lstm_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
base_training_configuration=base_cnn_lstm_training_config,
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"lstm_hidden_size": 32,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"lstm_hidden_size": 64,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
model_input_run = get_run_config(
|
||||
run_name="cnn_lstm_input_size_run",
|
||||
base_model_configuration=cnn_lstm_base_config |
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_cnn_lstm_training_config,
|
||||
model_config_kwargs_list=input_run_kwargs_list,
|
||||
)
|
||||
fertility_input_run = get_run_config(
|
||||
run_name="cnn_lstm_fertility_input_run",
|
||||
base_model_configuration=cnn_lstm_base_config |
|
||||
{
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_cnn_lstm_training_config,
|
||||
model_config_kwargs_list=[
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=10,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=80,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
)
|
||||
|
||||
],
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="cnn_lstm_complexity_run",
|
||||
base_model_configuration=cnn_lstm_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
},
|
||||
base_training_configuration=base_cnn_lstm_training_config,
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 16,
|
||||
"lstm_hidden_size": 16,
|
||||
"num_layers": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"lstm_hidden_size": 32,
|
||||
"num_layers": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"lstm_hidden_size": 32,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"lstm_hidden_size": 64,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 256,
|
||||
"lstm_hidden_size": 256,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"lstm_hidden_size": 512,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
best_config_run = get_run_config(
|
||||
run_name="cnn_lstm_best_config_run",
|
||||
base_model_configuration=cnn_lstm_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"lstm_hidden_size": 512,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_cnn_lstm_training_config | {
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-4,
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
},
|
||||
model_config_kwargs_list=[],
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
@@ -1,127 +0,0 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.feature_config import feature_config
|
||||
from models.cnn import CNNTransformer
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.lstm import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
take_every_nth = int(288 / 12)
|
||||
shift_in_hours = 12
|
||||
input_window_length = (288 // take_every_nth) * 80
|
||||
# output_window_length = (288 // take_every_nth) * 1
|
||||
output_window_length = 1
|
||||
output_window_offset = input_window_length + (288 // take_every_nth) * 0
|
||||
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
|
||||
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
|
||||
|
||||
max_lr = 1e-5
|
||||
batch_size = 256
|
||||
|
||||
|
||||
def get_cnn_run_config(
|
||||
run_name: str,
|
||||
run_description: str,
|
||||
input_window_length: int,
|
||||
output_window_length: int,
|
||||
take_every_nth: int,
|
||||
shift_in_hours: int,
|
||||
output_window_offset: int,
|
||||
batch_size: int,
|
||||
model_parameters: dict,
|
||||
max_lr: float,
|
||||
num_epochs: int,
|
||||
patience: int,
|
||||
feature_config: dict):
|
||||
"""
|
||||
Get the configuration for the CNN model
|
||||
Returns:
|
||||
run_configuration: configuration for the CNN model
|
||||
"""
|
||||
base_model_config = {
|
||||
"model_name": "cnn_regressor",
|
||||
"version": "1.0.0",
|
||||
"model_class": CNNTransformer,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"preprocessing": {
|
||||
"window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
|
||||
"take_every_nth": take_every_nth,
|
||||
"min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
|
||||
},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
"model_parameters": {
|
||||
**model_parameters | {"seq_len": input_window_length},
|
||||
},
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"output_window_offset": output_window_offset,
|
||||
}
|
||||
|
||||
base_training_config = {
|
||||
"batch_size": batch_size,
|
||||
"model_class": CNNTransformer,
|
||||
"learning_parameters": {
|
||||
"learning_rate": max_lr * (batch_size / 4),
|
||||
"epochs": num_epochs,
|
||||
"patience": patience
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
return {
|
||||
"name": run_name,
|
||||
"description": run_description,
|
||||
"model_configuration": base_model_config,
|
||||
"training_configuration": base_training_config,
|
||||
}
|
||||
|
||||
|
||||
run_configuration = {
|
||||
"runs": [
|
||||
get_cnn_run_config(
|
||||
run_name="cnn_ovulation_regression",
|
||||
run_description="CNN model for ovulation regression",
|
||||
input_window_length=input_window_length,
|
||||
output_window_length=output_window_length,
|
||||
take_every_nth=take_every_nth,
|
||||
shift_in_hours=shift_in_hours,
|
||||
output_window_offset=output_window_offset,
|
||||
batch_size=batch_size,
|
||||
model_parameters={
|
||||
"cnn_channels": 64,
|
||||
"kernel_size": 3,
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
max_lr=max_lr,
|
||||
num_epochs=10,
|
||||
patience=3,
|
||||
feature_config=feature_config,
|
||||
)
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.config_utils import *
|
||||
from configs.feature_config import feature_config
|
||||
from experiment_setup import get_fertility_based_eval_functions
|
||||
from models.cnn_transformer import CNNTransformer
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
cnn_transformer_base_config = {
|
||||
"model_name": "cnn_transformer",
|
||||
"version": "1.0.0",
|
||||
"model_class": CNNTransformer,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
}
|
||||
|
||||
cnn_transformer_base_training_config = {
|
||||
"batch_size": 128,
|
||||
"model_class": CNNTransformer,
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-5,
|
||||
"epochs": 30,
|
||||
"patience": 5
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="cnn_transformer_complexity_run",
|
||||
base_model_configuration=cnn_transformer_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
base_training_configuration=cnn_transformer_base_training_config,
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
model_input_run = get_run_config(
|
||||
run_name="cnn_transformer_input_size_run",
|
||||
base_model_configuration=cnn_transformer_base_config |
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=cnn_transformer_base_training_config,
|
||||
model_config_kwargs_list=input_run_kwargs_list,
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
fertility_input_run = get_run_config(
|
||||
run_name="cnn_transformer_fertility_input_run",
|
||||
base_model_configuration=cnn_transformer_base_config |
|
||||
{
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=cnn_transformer_base_training_config,
|
||||
model_config_kwargs_list=[
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=10,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=80,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])",
|
||||
}
|
||||
},
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="cnn_transformer_complexity_run",
|
||||
base_model_configuration=cnn_transformer_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
# "model_parameters": {
|
||||
# "embed_dim": 256,
|
||||
# "num_enc_layers": 2,
|
||||
# "num_heads": 2,
|
||||
# },
|
||||
},
|
||||
base_training_configuration=cnn_transformer_base_training_config,
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 16,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 256,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 8,
|
||||
"num_heads": 8,
|
||||
},
|
||||
},
|
||||
],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
best_config_run = get_run_config(
|
||||
run_name="cnn_transformer_best_config_run",
|
||||
base_model_configuration=cnn_transformer_base_config | get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) |
|
||||
{
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=cnn_transformer_base_training_config | {
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-4,
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
},
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=[],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -1,133 +0,0 @@
|
||||
from typing import Callable
|
||||
|
||||
import configs
|
||||
from configs.feature_config import feature_config
|
||||
from configs.cnn_run_config import get_cnn_run_config
|
||||
from configs.transformer_run_config import get_transformer_run_config
|
||||
|
||||
take_every_nth = int(288 / 12)
|
||||
shift_in_hours = 12
|
||||
input_window_length = (288 // take_every_nth) * 80
|
||||
# output_window_length = (288 // take_every_nth) * 1
|
||||
output_window_length = 1
|
||||
output_window_offset = input_window_length + (288 // take_every_nth) * 0
|
||||
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
|
||||
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
|
||||
|
||||
|
||||
def config_generator(config_gen_fn: Callable,
|
||||
fixed_params: dict,
|
||||
variable_param_configs: list,
|
||||
) -> list:
|
||||
configs = list()
|
||||
for variable_config in variable_param_configs:
|
||||
config = fixed_params.copy()
|
||||
for param_name, param_value in variable_config.items():
|
||||
# update the config with the variable parameter, values can be None, if default should be used
|
||||
if param_value is not None:
|
||||
config[param_name] = param_value
|
||||
configs.append(config_gen_fn(**config))
|
||||
|
||||
return configs
|
||||
|
||||
|
||||
run_configuration = {
|
||||
"name": "ovulation_regression",
|
||||
"runs": config_generator(
|
||||
get_cnn_run_config,
|
||||
fixed_params={
|
||||
"run_name": "cnn_ovulation_regression",
|
||||
"run_description": "CNN model for ovulation regression",
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"take_every_nth": take_every_nth,
|
||||
"shift_in_hours": shift_in_hours,
|
||||
"output_window_offset": output_window_offset,
|
||||
"batch_size": 128,
|
||||
"max_lr": 1e-5,
|
||||
"num_epochs": 10,
|
||||
"patience": 3,
|
||||
"feature_config": feature_config,
|
||||
},
|
||||
variable_param_configs=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"cnn_channels": 32,
|
||||
"kernel_size": 3,
|
||||
"embed_dim": 32,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"batch_size": 64,
|
||||
"model_parameters": {
|
||||
"cnn_channels": 64,
|
||||
"kernel_size": 3,
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"batch_size": 64,
|
||||
"model_parameters": {
|
||||
"cnn_channels": 128,
|
||||
"kernel_size": 5,
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
]
|
||||
) + config_generator(
|
||||
get_transformer_run_config,
|
||||
fixed_params={
|
||||
"run_name": "transformer_ovulation_regression",
|
||||
"run_description": "Transformer model for ovulation regression",
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"take_every_nth": take_every_nth,
|
||||
"shift_in_hours": shift_in_hours,
|
||||
"output_window_offset": output_window_offset,
|
||||
"batch_size": 128,
|
||||
"max_lr": 1e-5,
|
||||
"num_epochs": 10,
|
||||
"patience": 3,
|
||||
"feature_config": feature_config,
|
||||
},
|
||||
variable_param_configs=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"batch_size": 64,
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"batch_size": 64,
|
||||
"model_parameters": {
|
||||
"embed_dim": 256,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"batch_size": 32,
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,356 @@
|
||||
from copy import deepcopy
|
||||
from typing import Callable
|
||||
|
||||
from vsm_datascience_common import constants
|
||||
|
||||
|
||||
def recursive_dict_update(base_dict: dict, update_dict: dict) -> dict:
|
||||
base_dict = deepcopy(base_dict)
|
||||
for key, value in update_dict.items():
|
||||
if isinstance(value, dict) and key in base_dict:
|
||||
base_dict[key] = recursive_dict_update(base_dict[key], value)
|
||||
else:
|
||||
base_dict[key] = value
|
||||
|
||||
return base_dict
|
||||
|
||||
|
||||
def recursive_provider_dict_update(base_dict: dict,
|
||||
provider_fns: dict,
|
||||
context: dict = None) -> dict:
|
||||
"""
|
||||
Recursively updates a base dictionary with values from provider functions.
|
||||
|
||||
Provider functions get intermediate dictionary as input and return a value.
|
||||
Args:
|
||||
base_dict: base dictionary to update.
|
||||
provider_fns: provider functions to use for updating the dictionary.
|
||||
context: base dict to use as context in recursive calls, defaults to None.
|
||||
|
||||
Returns:
|
||||
A new dictionary with updated values from provider functions.
|
||||
|
||||
"""
|
||||
|
||||
if context is None:
|
||||
base_dict = deepcopy(base_dict)
|
||||
context = base_dict
|
||||
|
||||
for key, provider_fn in provider_fns.items():
|
||||
if isinstance(provider_fn, dict):
|
||||
# if the value is a dictionary, recursively update it
|
||||
base_dict[key] = recursive_provider_dict_update(base_dict.get(key, {}), provider_fn, context)
|
||||
elif callable(provider_fn):
|
||||
# call the provider function with the current base_dict as input
|
||||
base_dict[key] = provider_fn(context)
|
||||
else:
|
||||
# if value is string and starts with "eval(", evaluate it
|
||||
if isinstance(provider_fn, str) and provider_fn.startswith("eval(") and provider_fn.endswith(")"):
|
||||
# evaluate the string as a Python expression
|
||||
eval_fn = eval(provider_fn[5:-1])
|
||||
base_dict[key] = eval_fn(context)
|
||||
else:
|
||||
# if the value is not callable, just set it
|
||||
base_dict[key] = provider_fn
|
||||
|
||||
return base_dict
|
||||
|
||||
|
||||
def config_generator(config_gen_fn: Callable,
|
||||
fixed_params: dict,
|
||||
variable_param_configs: list,
|
||||
) -> list:
|
||||
configs = list()
|
||||
for variable_config in variable_param_configs:
|
||||
config = fixed_params.copy()
|
||||
for param_name, param_value in variable_config.items():
|
||||
# update the config with the variable parameter, values can be None, if default should be used
|
||||
if param_value is not None:
|
||||
config[param_name] = param_value
|
||||
configs.append(config_gen_fn(**config))
|
||||
|
||||
return configs
|
||||
|
||||
|
||||
def get_run_config(
|
||||
run_name: str,
|
||||
base_model_configuration: dict,
|
||||
base_training_configuration: dict,
|
||||
evaluation_configuration: dict = None,
|
||||
model_config_kwargs_list: list = None,
|
||||
training_config_kwargs_list: list = None,
|
||||
*args, **kwargs) -> dict:
|
||||
"""
|
||||
Generates a list of training configurations for a given model and training setup.
|
||||
|
||||
Args and kwargs are forwarded to the `get_training_config` function.
|
||||
Args:
|
||||
run_name: name of the run.
|
||||
base_model_configuration: base model configuration dictionary.
|
||||
base_training_configuration: base training configuration dictionary.
|
||||
evaluation_configuration: evaluation configuration dictionary, defaults to None.
|
||||
model_config_kwargs_list: list of dictionaries containing additional keyword arguments for model configuration.
|
||||
training_config_kwargs_list: list of dictionaries containing additional keyword arguments for training configuration.
|
||||
Returns:
|
||||
A dictionary containing the run configuration with a list of training configurations.
|
||||
|
||||
"""
|
||||
|
||||
if model_config_kwargs_list is None:
|
||||
model_config_kwargs_list = list()
|
||||
if training_config_kwargs_list is None:
|
||||
training_config_kwargs_list = list()
|
||||
|
||||
training_configs = list()
|
||||
num_trainings = max(len(model_config_kwargs_list), len(training_config_kwargs_list), 1)
|
||||
|
||||
num_model_kwargs = len(model_config_kwargs_list)
|
||||
num_training_kwargs = len(training_config_kwargs_list)
|
||||
|
||||
# if num_model_kwargs == 0 and num_training_kwargs == 0:
|
||||
# raise ValueError("Both model_config_kwargs_list and training_config_kwargs_list cannot be empty.")
|
||||
|
||||
for i in range(num_trainings):
|
||||
# list length can either be equal, one, or zero
|
||||
if num_model_kwargs == 0:
|
||||
model_config_kwargs = {}
|
||||
elif num_model_kwargs == 1:
|
||||
model_config_kwargs = model_config_kwargs_list[0]
|
||||
else:
|
||||
model_config_kwargs = model_config_kwargs_list[i]
|
||||
|
||||
if num_training_kwargs == 0:
|
||||
training_config_kwargs = {}
|
||||
elif num_training_kwargs == 1:
|
||||
training_config_kwargs = training_config_kwargs_list[0]
|
||||
else:
|
||||
training_config_kwargs = training_config_kwargs_list[i]
|
||||
|
||||
training_config = get_training_config(
|
||||
training_name=f"{run_name}_{i}",
|
||||
training_description=f"{run_name} training configuration {i}",
|
||||
base_model_config=base_model_configuration,
|
||||
model_config_kwargs=model_config_kwargs,
|
||||
base_training_config=base_training_configuration,
|
||||
training_config_kwargs=training_config_kwargs,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
training_configs.append(training_config)
|
||||
return {
|
||||
"name": run_name,
|
||||
"runs": training_configs,
|
||||
"evaluation_configuration": evaluation_configuration if evaluation_configuration is not None else {}
|
||||
}
|
||||
|
||||
|
||||
def get_training_config(
|
||||
training_name: str,
|
||||
training_description: str,
|
||||
base_model_config: dict,
|
||||
model_config_kwargs: dict,
|
||||
base_training_config: dict,
|
||||
training_config_kwargs: dict,
|
||||
additional_model_config_provider_fns: dict = None,
|
||||
additional_training_config_provider_fns: dict = None,
|
||||
*args, **kwargs) -> dict:
|
||||
"""
|
||||
Generates a training configuration for a given model and training setup.
|
||||
Args:
|
||||
training_name: Name of the training.
|
||||
training_description: Description of the training.
|
||||
base_model_config: Base model configuration dictionary.
|
||||
model_config_kwargs: Additional keyword arguments for model configuration.
|
||||
base_training_config: Base training configuration dictionary.
|
||||
training_config_kwargs: Additional keyword arguments for training configuration.
|
||||
additional_model_config_provider_fns: Additional provider functions for parts of the config that depend on the config itself
|
||||
such as input length as model parameter.
|
||||
additional_training_config_provider_fns: Additional provider functions for parts of the config that depend on the config itself
|
||||
such as input length as model parameter.
|
||||
Returns:
|
||||
A dictionary containing the training configuration with model and training parameters.
|
||||
"""
|
||||
|
||||
# recursively merge the base model configuration with the provided kwargs
|
||||
model_configuration = recursive_dict_update(base_model_config, model_config_kwargs)
|
||||
|
||||
# update based on provider functions
|
||||
if additional_model_config_provider_fns:
|
||||
model_configuration = recursive_provider_dict_update(model_configuration, additional_model_config_provider_fns)
|
||||
|
||||
required_model_configuration_keys = [
|
||||
"model_name",
|
||||
"version",
|
||||
"model_class",
|
||||
"feature_config",
|
||||
"preprocessing",
|
||||
"batch_fn",
|
||||
"collate_fn",
|
||||
"model_creation_fn",
|
||||
"model_save_fn",
|
||||
"model_load_fn",
|
||||
"batch_loss_fn",
|
||||
"actual_fn",
|
||||
"predict_fn",
|
||||
"model_parameters",
|
||||
"input_window_length",
|
||||
"output_window_length",
|
||||
"output_window_offset"
|
||||
]
|
||||
if not all(key in model_configuration for key in required_model_configuration_keys):
|
||||
missing_keys = [
|
||||
key for key in required_model_configuration_keys if key not in base_model_config
|
||||
]
|
||||
raise ValueError(f"Missing required keys in model configuration: {', '.join(missing_keys)}")
|
||||
|
||||
training_configuration = recursive_dict_update(base_training_config, training_config_kwargs)
|
||||
|
||||
# update based on provider functions
|
||||
if additional_training_config_provider_fns:
|
||||
training_configuration = recursive_provider_dict_update(training_configuration,
|
||||
additional_training_config_provider_fns)
|
||||
|
||||
# if "model_class" is not in training_configuration, set it to the model class from the model configuration
|
||||
if "model_class" not in training_configuration:
|
||||
training_configuration["model_class"] = base_model_config["model_class"]
|
||||
|
||||
required_training_configuration_keys = [
|
||||
"model_class",
|
||||
"batch_size",
|
||||
"learning_parameters",
|
||||
"max_grad_norm",
|
||||
"train_size",
|
||||
"val_size",
|
||||
"test_size"
|
||||
]
|
||||
|
||||
if not all(key in training_configuration for key in required_training_configuration_keys):
|
||||
missing_keys = [
|
||||
key for key in required_training_configuration_keys if key not in training_configuration
|
||||
]
|
||||
raise ValueError(f"Missing required keys in training configuration: {', '.join(missing_keys)}")
|
||||
|
||||
return {
|
||||
"name": training_name,
|
||||
"description": training_description,
|
||||
"model_configuration": model_configuration,
|
||||
"training_configuration": training_configuration,
|
||||
}
|
||||
|
||||
|
||||
def get_data_config_set(
|
||||
values_per_day: int,
|
||||
shift_in_hours: int,
|
||||
input_window_length_in_days: int,
|
||||
output_window_length: int,
|
||||
output_window_offset: int,
|
||||
) -> dict:
|
||||
measurements_per_day = constants.MEASUREMENTS_PER_DAY
|
||||
take_every_nth = int(measurements_per_day / values_per_day)
|
||||
input_window_length = (measurements_per_day // take_every_nth) * input_window_length_in_days
|
||||
return {
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"output_window_offset": input_window_length + output_window_offset,
|
||||
"preprocessing": {
|
||||
"window_shift": max(int((measurements_per_day // take_every_nth) / 24 * shift_in_hours), 1),
|
||||
"take_every_nth": take_every_nth,
|
||||
"min_input_length_fraction_for_padding":
|
||||
((measurements_per_day // take_every_nth) * 4) / input_window_length,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
input_run_kwargs_list = [
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=10,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=40,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=80,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=1,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=2,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=4,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=24,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=48,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=72,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
),
|
||||
get_data_config_set(
|
||||
values_per_day=288,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=20,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
)
|
||||
]
|
||||
@@ -7,6 +7,7 @@ from torch import nn
|
||||
import numpy as np
|
||||
|
||||
from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler
|
||||
from torch.nn import BCEWithLogitsLoss, MSELoss
|
||||
from torch.optim.lr_scheduler import OneCycleLR
|
||||
from torch.optim import AdamW
|
||||
|
||||
@@ -24,6 +25,10 @@ feature_config = {
|
||||
"$gt": constants.MEASUREMENTS_PER_DAY * 10,
|
||||
"$lt": constants.MEASUREMENTS_PER_DAY * 150
|
||||
},
|
||||
"duration_in_days": {
|
||||
"$gt": 10,
|
||||
"$lt": 150,
|
||||
},
|
||||
"ends_at": {
|
||||
"$exists": True,
|
||||
"$ne": None
|
||||
@@ -49,6 +54,21 @@ feature_config = {
|
||||
"static_categorical_features": [
|
||||
],
|
||||
"static_continuous_features": [
|
||||
{
|
||||
"name": "user_age",
|
||||
"fn": get_user_age,
|
||||
"scaler": RobustScaler
|
||||
},
|
||||
{
|
||||
"name": "user_age",
|
||||
"fn": get_user_height,
|
||||
"scaler": RobustScaler
|
||||
},
|
||||
{
|
||||
"name": "user_weight",
|
||||
"fn": get_user_weight,
|
||||
"scaler": RobustScaler
|
||||
},
|
||||
{
|
||||
"name": "average_cycle_length",
|
||||
"fn": get_cycle_length_stats,
|
||||
@@ -133,35 +153,40 @@ feature_config = {
|
||||
},
|
||||
],
|
||||
"target_features": [
|
||||
# {
|
||||
# "name": "fertility_probability",
|
||||
# "fn": partial(get_fertility_probability, shift=0),
|
||||
# "scaler": MinMaxScaler,
|
||||
# "accumulation_fn": np.max
|
||||
# },
|
||||
# {
|
||||
# "name": "ov_over_probability",
|
||||
# "fn": partial(get_ov_over_probability, shift=0),
|
||||
# "scaler": MinMaxScaler,
|
||||
# "accumulation_fn": np.max
|
||||
# },
|
||||
{
|
||||
"name": "fertility_probability",
|
||||
"fn": partial(get_fertility_probability, shift=0),
|
||||
"scaler": MinMaxScaler,
|
||||
"accumulation_fn": np.max,
|
||||
"loss_fn": MSELoss,
|
||||
},
|
||||
{
|
||||
"name": "ov_over_probability",
|
||||
"fn": partial(get_ov_over_probability, shift=0),
|
||||
"scaler": MinMaxScaler,
|
||||
"accumulation_fn": np.max,
|
||||
"loss_fn": BCEWithLogitsLoss,
|
||||
},
|
||||
{
|
||||
"name": "days_relative_to_ov",
|
||||
"fn": get_days_relative_to_ov,
|
||||
"scaler": RobustScaler,
|
||||
"accumulation_fn": np.max
|
||||
"accumulation_fn": np.max,
|
||||
"loss_fn": MSELoss,
|
||||
},
|
||||
{
|
||||
"name": "ov_day",
|
||||
"fn": get_ov_day,
|
||||
"scaler": RobustScaler,
|
||||
"accumulation_fn": np.max
|
||||
"accumulation_fn": np.max,
|
||||
"loss_fn": MSELoss,
|
||||
},
|
||||
{
|
||||
"name": "is_biphasic",
|
||||
"fn": get_is_biphasic,
|
||||
"scaler": MinMaxScaler,
|
||||
"accumulation_fn": np.max
|
||||
"accumulation_fn": np.max,
|
||||
"loss_fn": BCEWithLogitsLoss
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.config_utils import *
|
||||
from configs.feature_config import feature_config
|
||||
from experiment_setup import get_fertility_based_eval_functions
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.lstm import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
lstm_base_config = {
|
||||
"model_name": "lstm",
|
||||
"version": "1.0.0",
|
||||
"model_class": LSTM,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
}
|
||||
|
||||
base_lstm_training_config = {
|
||||
"batch_size": 256,
|
||||
"model_class": LSTM,
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-4,
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
|
||||
model_input_run = get_run_config(
|
||||
run_name="lstm_input_size_run",
|
||||
base_model_configuration=lstm_base_config |
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 64,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_lstm_training_config,
|
||||
model_config_kwargs_list=input_run_kwargs_list
|
||||
)
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="lstm_fertility_complexity_run",
|
||||
base_model_configuration=lstm_base_config | get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
# "model_parameters": {
|
||||
# "lstm_hidden_size": 128,
|
||||
# "num_layers": 4,
|
||||
# },
|
||||
},
|
||||
base_training_configuration=base_lstm_training_config,
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 16,
|
||||
"num_layers": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 32,
|
||||
"num_layers": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 32,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 64,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 256,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 512,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
],
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
fertility_input_run = get_run_config(
|
||||
run_name="lstm_fertility_input_size_run",
|
||||
base_model_configuration=lstm_base_config |
|
||||
{
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_lstm_training_config,
|
||||
model_config_kwargs_list=input_run_kwargs_list,
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
|
||||
best_config_run = get_run_config(
|
||||
run_name="lstm_best_config_run",
|
||||
base_model_configuration=lstm_base_config | get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"lstm_hidden_size": 512,
|
||||
"num_layers": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_lstm_training_config | {
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-4,
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": 50,
|
||||
"patience": 7,
|
||||
},
|
||||
},
|
||||
model_config_kwargs_list=[],
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
)
|
||||
@@ -1,127 +0,0 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.feature_config import feature_config
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.lstm import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
take_every_nth = int(288 / 12)
|
||||
shift_in_hours = 12
|
||||
input_window_length = (288 // take_every_nth) * 80
|
||||
# output_window_length = (288 // take_every_nth) * 1
|
||||
output_window_length = 1
|
||||
output_window_offset = input_window_length + (288 // take_every_nth) * 0
|
||||
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
|
||||
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
|
||||
|
||||
max_lr = 1e-5
|
||||
batch_size = 128
|
||||
|
||||
|
||||
def get_lstm_run_config(
|
||||
run_name: str,
|
||||
run_description: str,
|
||||
input_window_length: int,
|
||||
output_window_length: int,
|
||||
take_every_nth: int,
|
||||
shift_in_hours: int,
|
||||
output_window_offset: int,
|
||||
batch_size: int,
|
||||
model_parameters: dict,
|
||||
max_lr: float,
|
||||
num_epochs: int,
|
||||
patience: int,
|
||||
feature_config: dict):
|
||||
"""
|
||||
Get the configuration for the LSTM model
|
||||
Returns:
|
||||
run_configuration: configuration for the LSTM model
|
||||
"""
|
||||
base_model_config = {
|
||||
"model_name": "lstm_regressor",
|
||||
"version": "1.0.0",
|
||||
"model_class": LSTMModel,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"preprocessing": {
|
||||
"window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
|
||||
"take_every_nth": take_every_nth,
|
||||
"min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
|
||||
},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
"model_parameters": {
|
||||
**model_parameters,
|
||||
},
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"output_window_offset": output_window_offset,
|
||||
}
|
||||
base_training_config = {
|
||||
"batch_size": batch_size,
|
||||
"model_class": LSTMModel,
|
||||
"learning_parameters": {
|
||||
"learning_rate": max_lr * (batch_size / 4),
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": num_epochs,
|
||||
"patience": patience,
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
return {
|
||||
"name": run_name,
|
||||
"description": run_description,
|
||||
"model_configuration": base_model_config,
|
||||
"training_configuration": base_training_config,
|
||||
}
|
||||
|
||||
|
||||
run_configuration = {
|
||||
"item_limit": 100,
|
||||
"runs": [
|
||||
get_lstm_run_config(
|
||||
run_name="lstm_regressor",
|
||||
run_description="LSTM regressor for fertility prediction",
|
||||
input_window_length=input_window_length,
|
||||
output_window_length=output_window_length,
|
||||
take_every_nth=take_every_nth,
|
||||
shift_in_hours=shift_in_hours,
|
||||
output_window_offset=output_window_offset,
|
||||
batch_size=batch_size,
|
||||
model_parameters={
|
||||
"cnn_channels": 64,
|
||||
"cnn_kernel_size": 3,
|
||||
"embed_dim": 128,
|
||||
"lstm_hidden_size": 128,
|
||||
"num_layers": 4,
|
||||
},
|
||||
max_lr=max_lr,
|
||||
num_epochs=1000,
|
||||
patience=50,
|
||||
feature_config=feature_config
|
||||
)
|
||||
]
|
||||
}
|
||||
@@ -1,109 +0,0 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.feature_config import feature_config
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.third_party.patch_tst.models.PatchTST import Model as PatchTST
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
take_every_nth = int(288 / 12)
|
||||
shift_in_hours = 12
|
||||
input_window_length = (288 // take_every_nth) * 80
|
||||
# output_window_length = (288 // take_every_nth) * 1
|
||||
output_window_length = 1
|
||||
output_window_offset = input_window_length + (288 // take_every_nth) * 0
|
||||
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
|
||||
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
|
||||
|
||||
max_lr = 1e-5
|
||||
batch_size = 64
|
||||
|
||||
run_configuration = {
|
||||
"item_limit": 100,
|
||||
"runs": [
|
||||
{
|
||||
"name": "lstm_ovulation_regression",
|
||||
"description": "LSTM model for ovulation regression",
|
||||
"model_configuration": {
|
||||
"model_name": "patch_tst_regressor",
|
||||
"version": "1.0.0",
|
||||
"model_class": PatchTST,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"preprocessing": {
|
||||
"window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
|
||||
"take_every_nth": take_every_nth,
|
||||
"min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
|
||||
},
|
||||
"batch_fn": produce_window_batches,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
"model_parameters": {
|
||||
"configs": {
|
||||
# core
|
||||
"seq_len": input_window_length,
|
||||
"pred_len": output_window_length,
|
||||
"seq_pred": False,
|
||||
# model
|
||||
"e_layers": 4,
|
||||
"n_heads": 4,
|
||||
"d_model": 128,
|
||||
"d_ff": 128,
|
||||
"dropout": 0.2,
|
||||
"fc_dropout": 0.2,
|
||||
"head_dropout": 0.0,
|
||||
"individual": True,
|
||||
# patch
|
||||
# "patch_len": input_window_length,
|
||||
"patch_len": int(288 / take_every_nth),
|
||||
"stride": int(288 / take_every_nth / 2),
|
||||
"padding_patch": 0,
|
||||
# preprocessing
|
||||
"revin": False,
|
||||
"affine": False,
|
||||
"subtract_last": False,
|
||||
# decomp
|
||||
"decomposition": True,
|
||||
"kernel_size": 3,
|
||||
}
|
||||
},
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"output_window_offset": output_window_offset,
|
||||
},
|
||||
"training_configuration": {
|
||||
"batch_size": batch_size,
|
||||
"model_class": PatchTST,
|
||||
"learning_parameters": {
|
||||
# "learning_rate": base_lr,
|
||||
"learning_rate": max_lr * (batch_size / 4),
|
||||
"epochs": 10,
|
||||
"patience": 3,
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.config_utils import *
|
||||
from configs.feature_config import feature_config
|
||||
from experiment_setup import get_fertility_based_eval_functions
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.transformer import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
base_transformer_model_config = {
|
||||
"model_name": "transformer",
|
||||
"version": "1.0.0",
|
||||
"model_class": TransformerModel,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"batch_fn": produce_window_batches,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
}
|
||||
|
||||
base_transformer_training_config = {
|
||||
"batch_size": 128,
|
||||
"model_class": TransformerModel,
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-5,
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
|
||||
fertility_complexity_run = get_run_config(
|
||||
run_name="transformer_fertility_complexity_run",
|
||||
base_model_configuration=base_transformer_model_config | get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
# "model_parameters": {
|
||||
# "embed_dim": 256,
|
||||
# "num_enc_layers": 2,
|
||||
# "num_heads": 2,
|
||||
# },
|
||||
},
|
||||
base_training_configuration=base_transformer_training_config,
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=[
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 16,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 32,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 1,
|
||||
"num_heads": 1,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 64,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 128,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 256,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 8,
|
||||
"num_heads": 8,
|
||||
},
|
||||
},
|
||||
],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
fertility_input_run = get_run_config(
|
||||
run_name="transformer_fertility_input_run",
|
||||
base_model_configuration=base_transformer_model_config | {
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 256,
|
||||
"num_enc_layers": 2,
|
||||
"num_heads": 2,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_transformer_training_config,
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=input_run_kwargs_list,
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
best_config_run = get_run_config(
|
||||
run_name="transformer_best_config_run",
|
||||
base_model_configuration=base_transformer_model_config | get_data_config_set(
|
||||
values_per_day=12,
|
||||
shift_in_hours=12,
|
||||
input_window_length_in_days=160,
|
||||
output_window_length=1,
|
||||
output_window_offset=0
|
||||
) |
|
||||
{
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
# "fertility_probability",
|
||||
# "ov_over_probability",
|
||||
"days_relative_to_ov",
|
||||
"is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"model_parameters": {
|
||||
"embed_dim": 512,
|
||||
"num_enc_layers": 4,
|
||||
"num_heads": 4,
|
||||
},
|
||||
},
|
||||
base_training_configuration=base_transformer_training_config | {
|
||||
"learning_parameters": {
|
||||
"max_lr": 1e-4,
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": 30,
|
||||
"patience": 5,
|
||||
},
|
||||
},
|
||||
evaluation_configuration={
|
||||
"eval_functions_getter": get_fertility_based_eval_functions
|
||||
},
|
||||
model_config_kwargs_list=[],
|
||||
# add model config provider function to set seq_len in model parameters
|
||||
additional_model_config_provider_fns={
|
||||
"model_parameters": {
|
||||
"seq_len": "eval(lambda x: x['input_window_length'])"
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -1,123 +0,0 @@
|
||||
from functools import partial
|
||||
|
||||
from configs.feature_config import feature_config
|
||||
|
||||
from utils.data_utils import *
|
||||
from utils.training import *
|
||||
|
||||
from models.transformer import *
|
||||
from models.utils import *
|
||||
from models.collation import *
|
||||
|
||||
take_every_nth = int(288 / 12)
|
||||
shift_in_hours = 12
|
||||
input_window_length = (288 // take_every_nth) * 80
|
||||
# output_window_length = (288 // take_every_nth) * 1
|
||||
output_window_length = 1
|
||||
output_window_offset = input_window_length + (288 // take_every_nth) * 0
|
||||
window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
|
||||
min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
|
||||
|
||||
transformer_batch_size = 128
|
||||
|
||||
|
||||
def get_transformer_run_config(
|
||||
run_name: str,
|
||||
run_description: str,
|
||||
input_window_length: int,
|
||||
output_window_length: int,
|
||||
take_every_nth: int,
|
||||
shift_in_hours: int,
|
||||
output_window_offset: int,
|
||||
batch_size: int,
|
||||
model_parameters: dict,
|
||||
max_lr: float,
|
||||
num_epochs: int,
|
||||
patience: int,
|
||||
feature_config: dict):
|
||||
"""
|
||||
Get the configuration for the Transformer model
|
||||
Returns:
|
||||
run_configuration: configuration for the Transformer model
|
||||
"""
|
||||
base_model_config = {
|
||||
"model_name": "transformer_regressor",
|
||||
"version": "1.0.0",
|
||||
"model_class": TransformerModel,
|
||||
"feature_config": feature_config | {"ignored_features":
|
||||
[
|
||||
"fertility_probability",
|
||||
"ov_over_probability",
|
||||
# "days_relative_to_ov",
|
||||
# "is_biphasic",
|
||||
"ov_day",
|
||||
]},
|
||||
"preprocessing": {
|
||||
"window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
|
||||
"take_every_nth": take_every_nth,
|
||||
"min_input_length_fraction_for_padding": min_input_length_fraction_for_padding,
|
||||
},
|
||||
"batch_fn": produce_window_batches,
|
||||
"model_creation_fn": simple_model_creation,
|
||||
"collate_fn": simple_x_y_collate,
|
||||
"model_save_fn": simple_model_save,
|
||||
"model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
|
||||
"batch_loss_fn": get_model_loss,
|
||||
"actual_fn": simple_get_y,
|
||||
"predict_fn": simple_x_y_predict,
|
||||
"model_parameters": {
|
||||
**model_parameters | {"seq_len": input_window_length},
|
||||
},
|
||||
"input_window_length": input_window_length,
|
||||
"output_window_length": output_window_length,
|
||||
"output_window_offset": output_window_offset,
|
||||
}
|
||||
base_training_config = {
|
||||
"batch_size": batch_size,
|
||||
"model_class": TransformerModel,
|
||||
"learning_parameters": {
|
||||
"learning_rate": max_lr * (batch_size / 4),
|
||||
# "learning_rate": base_lr * (batch_size / 4),
|
||||
"epochs": num_epochs,
|
||||
"patience": patience,
|
||||
},
|
||||
"loss_functions": [
|
||||
nn.MSELoss(),
|
||||
nn.BCEWithLogitsLoss(),
|
||||
],
|
||||
"max_grad_norm": 1.0,
|
||||
"train_size": 0.7,
|
||||
"val_size": 0.15,
|
||||
"test_size": 0.15,
|
||||
}
|
||||
return {
|
||||
"name": run_name,
|
||||
"description": run_description,
|
||||
"model_configuration": base_model_config,
|
||||
"training_configuration": base_training_config
|
||||
}
|
||||
|
||||
|
||||
run_configuration = {
|
||||
"runs": [
|
||||
get_transformer_run_config(
|
||||
run_name="transformer_ovulation_regression",
|
||||
run_description="Transformer regressor for ovulation prediction",
|
||||
input_window_length=input_window_length,
|
||||
output_window_length=output_window_length,
|
||||
take_every_nth=take_every_nth,
|
||||
shift_in_hours=shift_in_hours,
|
||||
output_window_offset=output_window_offset,
|
||||
batch_size=transformer_batch_size,
|
||||
model_parameters={
|
||||
"embed_dim": 128,
|
||||
"num_heads": 4,
|
||||
"num_enc_layers": 4,
|
||||
},
|
||||
max_lr=1e-5,
|
||||
num_epochs=20,
|
||||
patience=3,
|
||||
feature_config=feature_config,
|
||||
)
|
||||
]
|
||||
}
|
||||
@@ -14,31 +14,26 @@ def get_eval_functions(model_configuration: dict):
|
||||
eval_functions = [
|
||||
{
|
||||
"name": "mean_absolute_error_overall",
|
||||
"input_index": 0,
|
||||
"eval_fn": sklearn.metrics.mean_absolute_error,
|
||||
"accumulation_fn": np.mean,
|
||||
},
|
||||
{
|
||||
"name": "mean_absolute_error_pre_ov",
|
||||
"input_index": 0,
|
||||
"eval_fn": pre_ov_error,
|
||||
"accumulation_fn": np.mean,
|
||||
},
|
||||
{
|
||||
"name": "mean_absolute_error_after_ov",
|
||||
"input_index": 0,
|
||||
"eval_fn": after_ov_error,
|
||||
"accumulation_fn": np.mean,
|
||||
},
|
||||
{
|
||||
"name": "mean_absolute_error_ov_in_days",
|
||||
"input_index": 0,
|
||||
"eval_fn": partial(ov_error, model_configuration=model_configuration),
|
||||
"accumulation_fn": np.mean,
|
||||
},
|
||||
{
|
||||
"name": "mean_absolute_error_five_days_before_ov",
|
||||
"input_index": 0,
|
||||
"eval_fn": partial(day_relative_to_ov_error,
|
||||
day_relative_to_ov=-5,
|
||||
model_configuration=model_configuration),
|
||||
@@ -46,3 +41,82 @@ def get_eval_functions(model_configuration: dict):
|
||||
}
|
||||
]
|
||||
return eval_functions
|
||||
|
||||
|
||||
def get_fertility_based_eval_functions(model_configuration: dict):
|
||||
eval_functions = [
|
||||
{
|
||||
"name": "mae_ov_over",
|
||||
"eval_fn": sklearn.metrics.mean_absolute_error,
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mae_ov_over_before_ov",
|
||||
"eval_fn": partial(get_ov_over_pre_ov_error, error_fn=sklearn.metrics.mean_absolute_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mae_ov_over_after_ov",
|
||||
"eval_fn": partial(get_ov_over_post_ov_error, error_fn=sklearn.metrics.mean_absolute_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mse_ov_over",
|
||||
"eval_fn": sklearn.metrics.mean_squared_error,
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mse_ov_over_before_ov",
|
||||
"eval_fn": partial(get_ov_over_pre_ov_error, error_fn=sklearn.metrics.mean_squared_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mse_ov_over_after_ov",
|
||||
"eval_fn": partial(get_ov_over_post_ov_error, error_fn=sklearn.metrics.mean_squared_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 1
|
||||
},
|
||||
{
|
||||
"name": "mae_fertility",
|
||||
"eval_fn": sklearn.metrics.mean_absolute_error,
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
{
|
||||
"name": "mae_during_fertility",
|
||||
"eval_fn": partial(get_during_fertility_error, error_fn=sklearn.metrics.mean_absolute_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
{
|
||||
"name": "mae_non_fertility",
|
||||
"eval_fn": partial(get_non_fertility_error, error_fn=sklearn.metrics.mean_absolute_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
{
|
||||
"name": "mse_fertility",
|
||||
"eval_fn": sklearn.metrics.mean_squared_error,
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
{
|
||||
"name": "mse_during_fertility",
|
||||
"eval_fn": partial(get_during_fertility_error, error_fn=sklearn.metrics.mean_squared_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
{
|
||||
"name": "mse_non_fertility",
|
||||
"eval_fn": partial(get_non_fertility_error, error_fn=sklearn.metrics.mean_squared_error),
|
||||
"accumulation_fn": np.mean,
|
||||
"input_index": 0
|
||||
},
|
||||
]
|
||||
|
||||
return eval_functions
|
||||
|
||||
@@ -0,0 +1,173 @@
|
||||
from tqdm import tqdm
|
||||
from vsm_datascience_common.cycle_database_connection.cycle_data import *
|
||||
from vsm_datascience_common.cycles.sequences import *
|
||||
|
||||
from utils.evaluation import aggregate_errors
|
||||
from utils.cycle_utils import get_ovulation_day
|
||||
from utils.feature_functions import *
|
||||
from utils.utils import recursive_dict_update
|
||||
|
||||
|
||||
def get_cycle_fertility_curve(cycle: dict, ov_day: int) -> np.ndarray:
|
||||
timestamps = get_timestamps(cycle)
|
||||
cycle_length = len(timestamps)
|
||||
|
||||
if ov_day is None:
|
||||
ov_index = None
|
||||
else:
|
||||
ov_timestamp = cycle["starts_at"] + timedelta(days=ov_day)
|
||||
ov_index = np.searchsorted(timestamps, ov_timestamp)
|
||||
|
||||
if ov_index is None:
|
||||
fertility_curve = np.full(cycle_length, 0.0)
|
||||
else:
|
||||
fertility_curve = get_fertility_curve(
|
||||
cycle_length,
|
||||
ov_index,
|
||||
get_fertility_probability_base_curve()[0],
|
||||
0
|
||||
)
|
||||
|
||||
return fertility_curve
|
||||
|
||||
|
||||
def get_ov_over_curve(cycle: dict, ov_day: int) -> np.ndarray:
|
||||
timestamps = get_timestamps(cycle)
|
||||
cycle_length = len(timestamps)
|
||||
|
||||
if ov_day is None:
|
||||
ov_index = None
|
||||
else:
|
||||
ov_timestamp = cycle["starts_at"] + timedelta(days=ov_day)
|
||||
ov_index = np.searchsorted(timestamps, ov_timestamp)
|
||||
|
||||
ov_over_curve = np.full(cycle_length, 0.0)
|
||||
if ov_index is not None:
|
||||
ov_over_curve[min(ov_index, cycle_length - 1):] = 1.0
|
||||
|
||||
return ov_over_curve
|
||||
|
||||
|
||||
def get_last_cycle_baseline(cycle: dict) -> tuple[np.ndarray, np.ndarray]:
|
||||
try:
|
||||
last_cycle = get_previous_cycle(cycle["_id"])
|
||||
except ValueError:
|
||||
last_cycle = None
|
||||
|
||||
current_ov_day = get_ovulation_day(cycle)
|
||||
if last_cycle is None:
|
||||
last_ov_day = 18
|
||||
else:
|
||||
last_ov_day = get_ovulation_day(last_cycle)
|
||||
|
||||
current_fertility_curve = get_cycle_fertility_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_ov_over_curve = get_ov_over_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_actual = np.concatenate((current_fertility_curve, current_ov_over_curve), axis=1)
|
||||
|
||||
last_fertility_curve = get_cycle_fertility_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_ov_over_curve = get_ov_over_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_actual = np.concatenate((last_fertility_curve, last_ov_over_curve), axis=1)
|
||||
|
||||
return last_actual, current_actual
|
||||
|
||||
|
||||
def get_population_mean_baseline(cycle: dict) -> tuple[np.ndarray, np.ndarray]:
|
||||
current_ov_day = get_ovulation_day(cycle)
|
||||
last_ov_day = 18
|
||||
|
||||
current_fertility_curve = get_cycle_fertility_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_ov_over_curve = get_ov_over_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_actual = np.concatenate((current_fertility_curve, current_ov_over_curve), axis=1)
|
||||
|
||||
last_fertility_curve = get_cycle_fertility_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_ov_over_curve = get_ov_over_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_actual = np.concatenate((last_fertility_curve, last_ov_over_curve), axis=1)
|
||||
|
||||
return last_actual, current_actual
|
||||
|
||||
|
||||
def get_user_mean_baseline(cycle: dict) -> tuple[np.ndarray, np.ndarray]:
|
||||
current_ov_day = get_ovulation_day(cycle)
|
||||
|
||||
previous_cycle = get_previous_cycles(cycle["_id"], 100)
|
||||
if previous_cycle is None:
|
||||
last_ov_day = 18
|
||||
else:
|
||||
previous_ovs = list()
|
||||
for prev_cycle in previous_cycle:
|
||||
ov_day = get_ovulation_day(prev_cycle)
|
||||
if ov_day is not None:
|
||||
previous_ovs.append(ov_day)
|
||||
if len(previous_ovs) == 0:
|
||||
last_ov_day = 18
|
||||
else:
|
||||
last_ov_day = int(np.mean(previous_ovs))
|
||||
|
||||
current_fertility_curve = get_cycle_fertility_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_ov_over_curve = get_ov_over_curve(cycle, current_ov_day).reshape(-1, 1)
|
||||
current_actual = np.concatenate((current_fertility_curve, current_ov_over_curve), axis=1)
|
||||
|
||||
last_fertility_curve = get_cycle_fertility_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_ov_over_curve = get_ov_over_curve(cycle, last_ov_day).reshape(-1, 1)
|
||||
last_actual = np.concatenate((last_fertility_curve, last_ov_over_curve), axis=1)
|
||||
|
||||
return last_actual, current_actual
|
||||
|
||||
|
||||
def get_baseline_evaluation_for_cycle(cycle: dict,
|
||||
cycle_number: int,
|
||||
predictor_fn: Callable,
|
||||
eval_fns: list) -> dict:
|
||||
current_predictions, current_actuals = predictor_fn(cycle)
|
||||
errors = dict()
|
||||
for eval_fn in eval_fns:
|
||||
if eval_fn is not None:
|
||||
eval_fn_name = eval_fn["name"]
|
||||
if eval_fn_name not in errors:
|
||||
errors[eval_fn_name] = dict()
|
||||
eval_function = eval_fn["eval_fn"]
|
||||
eval_fn_indices = [eval_fn["input_index"]] if "input_index" in eval_fn else range(
|
||||
len(current_predictions[0]))
|
||||
|
||||
for output_index in eval_fn_indices:
|
||||
if any(np.isnan(current_predictions[:, output_index])):
|
||||
continue
|
||||
|
||||
input_preds = current_predictions[:, output_index]
|
||||
input_actuals = current_actuals[:, output_index]
|
||||
error = eval_function(input_preds, input_actuals)
|
||||
|
||||
if np.isnan(error):
|
||||
# skip if error is nan
|
||||
continue
|
||||
|
||||
if f"after_{cycle_number}" not in errors[eval_fn_name]:
|
||||
errors[eval_fn_name][f"after_{cycle_number}"] = dict()
|
||||
|
||||
if output_index not in errors[eval_fn_name][f"after_{cycle_number}"]:
|
||||
errors[eval_fn_name][f"after_{cycle_number}"][output_index] = list()
|
||||
|
||||
errors[eval_fn_name][f"after_{cycle_number}"][output_index].append(error)
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
def get_errors_for_users(user_ids: list,
|
||||
key_stats: dict,
|
||||
predictor_fn: Callable,
|
||||
eval_fns: list,
|
||||
aggregate: bool = True,
|
||||
show_progress: bool = False) -> dict:
|
||||
errors = dict()
|
||||
for key in tqdm(user_ids, disable=not show_progress):
|
||||
user_cycle_ids = [c["cycle_id"] for c in key_stats["by_key"][key]["cycle_stats"]]
|
||||
user_cycles = [get_cycle_by_id(cycle_id) for cycle_id in user_cycle_ids]
|
||||
# sort by starts_at
|
||||
user_cycles = sorted(user_cycles, key=lambda c: c["starts_at"])
|
||||
for i, cycle in enumerate(user_cycles):
|
||||
current_errors = get_baseline_evaluation_for_cycle(cycle, i, predictor_fn, eval_fns)
|
||||
errors = recursive_dict_update(errors, current_errors)
|
||||
|
||||
if aggregate:
|
||||
errors = aggregate_errors(errors, eval_fns)
|
||||
return errors
|
||||
@@ -0,0 +1,52 @@
|
||||
from torch import nn
|
||||
|
||||
|
||||
class CNNLSTM(nn.Module):
|
||||
def __init__(self,
|
||||
input_dim: int,
|
||||
output_dim: int,
|
||||
embed_dim: int,
|
||||
lstm_hidden_size=64,
|
||||
num_layers=1):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=9, stride=2), # 288 → ~140
|
||||
nn.ReLU(),
|
||||
nn.AdaptiveAvgPool1d(output_size=128), # force to 128
|
||||
nn.Conv1d(input_dim, input_dim, kernel_size=5, stride=2), # 128 → ~62
|
||||
nn.ReLU(),
|
||||
nn.AdaptiveAvgPool1d(output_size=48), # final fixed length
|
||||
)
|
||||
|
||||
# linear projection to embed dim
|
||||
self.input_proj = nn.Linear(input_dim, embed_dim)
|
||||
|
||||
self.lstm = nn.LSTM(
|
||||
input_size=embed_dim,
|
||||
hidden_size=lstm_hidden_size,
|
||||
num_layers=num_layers,
|
||||
batch_first=True,
|
||||
dropout=0.4,
|
||||
)
|
||||
self.head = nn.Sequential(
|
||||
nn.Linear(lstm_hidden_size, output_dim) # Output is a scalar Δt
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
# x: (batch_size, seq_len, input_size)
|
||||
x = x.permute(0, 2, 1)
|
||||
# x: (batch_size, input_size, seq_len)
|
||||
x = self.conv(x)
|
||||
# x: (batch_size, embed_dim, seq_len)
|
||||
x = x.permute(0, 2, 1)
|
||||
# x: (batch_size, seq_len, embed_dim)
|
||||
# project to embed dim
|
||||
x = self.input_proj(x)
|
||||
# x: (batch_size, seq_len, embed_dim)
|
||||
x, _ = self.lstm(x)
|
||||
# x: (batch_size, seq_len, lstm_hidden_size)
|
||||
x = x[:, -1, :] # Get the last time step
|
||||
# x: (batch_size, lstm_hidden_size)
|
||||
x = self.head(x)
|
||||
# x: (batch_size, output_size)
|
||||
return x
|
||||
@@ -9,26 +9,27 @@ class CNNTransformer(nn.Module):
|
||||
input_dim,
|
||||
output_dim,
|
||||
seq_len,
|
||||
cnn_channels,
|
||||
kernel_size,
|
||||
embed_dim,
|
||||
num_enc_layers,
|
||||
num_heads):
|
||||
super().__init__()
|
||||
self.cnn = nn.Sequential(
|
||||
nn.Conv1d(input_dim, cnn_channels, kernel_size=kernel_size, padding=1),
|
||||
nn.BatchNorm1d(cnn_channels),
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=9, stride=2), # 288 → ~140
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(cnn_channels, embed_dim, kernel_size=kernel_size, padding=1),
|
||||
nn.BatchNorm1d(embed_dim),
|
||||
nn.ReLU()
|
||||
nn.AdaptiveAvgPool1d(output_size=128), # force to 128
|
||||
nn.Conv1d(input_dim, input_dim, kernel_size=5, stride=2), # 128 → ~62
|
||||
nn.ReLU(),
|
||||
nn.AdaptiveAvgPool1d(output_size=48), # final fixed length
|
||||
)
|
||||
|
||||
# linear projection to embed dim
|
||||
self.input_proj = nn.Linear(input_dim, embed_dim)
|
||||
|
||||
# Compute positional embedding ONCE at init
|
||||
pe = self._get_sinusoidal_embedding(seq_len, embed_dim) # (seq_len, embed_dim)
|
||||
self.register_buffer('pos_embed', pe.unsqueeze(0)) # (1, seq_len, embed_dim)
|
||||
|
||||
encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
|
||||
encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads, dropout=0.1)
|
||||
self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers)
|
||||
self.pool = nn.AdaptiveAvgPool1d(1)
|
||||
self.head = nn.Linear(embed_dim, output_dim)
|
||||
@@ -42,17 +43,21 @@ class CNNTransformer(nn.Module):
|
||||
return pe # (seq_len, embed_dim)
|
||||
|
||||
def forward(self, x):
|
||||
# x: (batch, seq_len, input_dim)
|
||||
x = x.permute(0, 2, 1) # (B, input_dim, seq_len)
|
||||
cnn_out = self.cnn(x) # (B, embed_dim, seq_len)
|
||||
cnn_out = cnn_out.permute(2, 0, 1) # (S, B, E) for Transformer
|
||||
# x: (B, seq_len, input_dim)
|
||||
x = x.permute(0, 2, 1) # → (B, input_dim, seq_len)
|
||||
cnn_out = self.conv(x) # → (B, input_dim, 48)
|
||||
|
||||
cnn_out = cnn_out.transpose(1, 2) # → (B, 48, input_dim)
|
||||
transformer_in = self.input_proj(cnn_out) # → (B, 48, embed_dim)
|
||||
|
||||
transformer_in = transformer_in.transpose(0, 1) # → (48, B, embed_dim)
|
||||
|
||||
# Add positional embedding
|
||||
pos_embed = self.pos_embed[:, :cnn_out.size(0), :] # (1, seq_len, embed_dim)
|
||||
pos_embed = self.pos_embed[:, :transformer_in.size(0), :] # (1, seq_len, embed_dim)
|
||||
pos_embed = pos_embed.transpose(0, 1) # → (seq_len, 1, embed_dim)
|
||||
cnn_out = cnn_out + pos_embed # broadcast over batch
|
||||
transformer_in = transformer_in + pos_embed # broadcast over batch
|
||||
|
||||
# Apply Transformer encoder
|
||||
enc = self.encoder(cnn_out) # (S, B, E)
|
||||
enc = self.encoder(transformer_in) # (S, B, E)
|
||||
pooled = enc.mean(0) # (B, E)
|
||||
return self.head(pooled)
|
||||
+7
-20
@@ -1,29 +1,22 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class LSTMModel(nn.Module):
|
||||
class LSTM(nn.Module):
|
||||
def __init__(self,
|
||||
input_dim: int,
|
||||
output_dim: int,
|
||||
cnn_channels: int,
|
||||
cnn_kernel_size: int,
|
||||
embed_dim: int,
|
||||
lstm_hidden_size=64,
|
||||
num_layers=1):
|
||||
num_layers=1,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
self.cnn = nn.Sequential(
|
||||
nn.Conv1d(input_dim, cnn_channels, kernel_size=cnn_kernel_size, padding=1),
|
||||
nn.BatchNorm1d(cnn_channels),
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(cnn_channels, embed_dim, kernel_size=cnn_kernel_size, padding=1),
|
||||
nn.BatchNorm1d(embed_dim),
|
||||
nn.ReLU()
|
||||
)
|
||||
self.lstm = nn.LSTM(
|
||||
input_size=embed_dim,
|
||||
input_size=input_dim,
|
||||
hidden_size=lstm_hidden_size,
|
||||
num_layers=num_layers,
|
||||
batch_first=True,
|
||||
dropout=0.5,
|
||||
**kwargs
|
||||
)
|
||||
self.head = nn.Sequential(
|
||||
nn.Linear(lstm_hidden_size, output_dim) # Output is a scalar Δt
|
||||
@@ -31,12 +24,6 @@ class LSTMModel(nn.Module):
|
||||
|
||||
def forward(self, x):
|
||||
# x: (batch_size, seq_len, input_size)
|
||||
x = x.permute(0, 2, 1)
|
||||
# x: (batch_size, input_size, seq_len)
|
||||
x = self.cnn(x)
|
||||
# x: (batch_size, embed_dim, seq_len)
|
||||
x = x.permute(0, 2, 1)
|
||||
# x: (batch_size, seq_len, embed_dim)
|
||||
x, _ = self.lstm(x)
|
||||
# x: (batch_size, seq_len, lstm_hidden_size)
|
||||
x = x[:, -1, :] # Get the last time step
|
||||
|
||||
@@ -1,164 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import math
|
||||
from math import sqrt
|
||||
import os
|
||||
|
||||
|
||||
class AutoCorrelation(nn.Module):
|
||||
"""
|
||||
AutoCorrelation Mechanism with the following two phases:
|
||||
(1) period-based dependencies discovery
|
||||
(2) time delay aggregation
|
||||
This block can replace the self-attention family mechanism seamlessly.
|
||||
"""
|
||||
def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(AutoCorrelation, self).__init__()
|
||||
self.factor = factor
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def time_delay_agg_training(self, values, corr):
|
||||
"""
|
||||
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
||||
This is for the training phase.
|
||||
"""
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
||||
index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
|
||||
weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
pattern = torch.roll(tmp_values, -int(index[i]), -1)
|
||||
delays_agg = delays_agg + pattern * \
|
||||
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
||||
return delays_agg
|
||||
|
||||
def time_delay_agg_inference(self, values, corr):
|
||||
"""
|
||||
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
||||
This is for the inference phase.
|
||||
"""
|
||||
batch = values.shape[0]
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# index init
|
||||
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
||||
weights = torch.topk(mean_value, top_k, dim=-1)[0]
|
||||
delay = torch.topk(mean_value, top_k, dim=-1)[1]
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values.repeat(1, 1, 1, 2)
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
|
||||
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
||||
delays_agg = delays_agg + pattern * \
|
||||
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
||||
return delays_agg
|
||||
|
||||
def time_delay_agg_full(self, values, corr):
|
||||
"""
|
||||
Standard version of Autocorrelation
|
||||
"""
|
||||
batch = values.shape[0]
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# index init
|
||||
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
weights = torch.topk(corr, top_k, dim=-1)[0]
|
||||
delay = torch.topk(corr, top_k, dim=-1)[1]
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values.repeat(1, 1, 1, 2)
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
tmp_delay = init_index + delay[..., i].unsqueeze(-1)
|
||||
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
||||
delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
|
||||
return delays_agg
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
if L > S:
|
||||
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
|
||||
values = torch.cat([values, zeros], dim=1)
|
||||
keys = torch.cat([keys, zeros], dim=1)
|
||||
else:
|
||||
values = values[:, :L, :, :]
|
||||
keys = keys[:, :L, :, :]
|
||||
|
||||
# period-based dependencies
|
||||
q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1)
|
||||
k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
|
||||
res = q_fft * torch.conj(k_fft)
|
||||
corr = torch.fft.irfft(res, dim=-1)
|
||||
|
||||
# time delay agg
|
||||
if self.training:
|
||||
V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
||||
else:
|
||||
V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
||||
|
||||
if self.output_attention:
|
||||
return (V.contiguous(), corr.permute(0, 3, 1, 2))
|
||||
else:
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class AutoCorrelationLayer(nn.Module):
|
||||
def __init__(self, correlation, d_model, n_heads, d_keys=None,
|
||||
d_values=None):
|
||||
super(AutoCorrelationLayer, self).__init__()
|
||||
|
||||
d_keys = d_keys or (d_model // n_heads)
|
||||
d_values = d_values or (d_model // n_heads)
|
||||
|
||||
self.inner_correlation = correlation
|
||||
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.value_projection = nn.Linear(d_model, d_values * n_heads)
|
||||
self.out_projection = nn.Linear(d_values * n_heads, d_model)
|
||||
self.n_heads = n_heads
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, _ = queries.shape
|
||||
_, S, _ = keys.shape
|
||||
H = self.n_heads
|
||||
|
||||
queries = self.query_projection(queries).view(B, L, H, -1)
|
||||
keys = self.key_projection(keys).view(B, S, H, -1)
|
||||
values = self.value_projection(values).view(B, S, H, -1)
|
||||
|
||||
out, attn = self.inner_correlation(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
attn_mask
|
||||
)
|
||||
out = out.view(B, L, -1)
|
||||
|
||||
return self.out_projection(out), attn
|
||||
@@ -1,173 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class my_Layernorm(nn.Module):
|
||||
"""
|
||||
Special designed layernorm for the seasonal part
|
||||
"""
|
||||
def __init__(self, channels):
|
||||
super(my_Layernorm, self).__init__()
|
||||
self.layernorm = nn.LayerNorm(channels)
|
||||
|
||||
def forward(self, x):
|
||||
x_hat = self.layernorm(x)
|
||||
bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
|
||||
return x_hat - bias
|
||||
|
||||
|
||||
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 EncoderLayer(nn.Module):
|
||||
"""
|
||||
Autoformer encoder layer with the progressive decomposition architecture
|
||||
"""
|
||||
def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
|
||||
super(EncoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.attention = attention
|
||||
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
|
||||
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
|
||||
self.decomp1 = series_decomp(moving_avg)
|
||||
self.decomp2 = series_decomp(moving_avg)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, attn_mask=None):
|
||||
new_x, attn = self.attention(
|
||||
x, x, x,
|
||||
attn_mask=attn_mask
|
||||
)
|
||||
x = x + self.dropout(new_x)
|
||||
x, _ = self.decomp1(x)
|
||||
y = x
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
res, _ = self.decomp2(x + y)
|
||||
return res, attn
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
"""
|
||||
Autoformer encoder
|
||||
"""
|
||||
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):
|
||||
attns = []
|
||||
if self.conv_layers is not None:
|
||||
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
x = conv_layer(x)
|
||||
attns.append(attn)
|
||||
x, attn = self.attn_layers[-1](x)
|
||||
attns.append(attn)
|
||||
else:
|
||||
for attn_layer in self.attn_layers:
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
attns.append(attn)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
return x, attns
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
"""
|
||||
Autoformer decoder layer with the progressive decomposition architecture
|
||||
"""
|
||||
def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
|
||||
moving_avg=25, dropout=0.1, activation="relu"):
|
||||
super(DecoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.self_attention = self_attention
|
||||
self.cross_attention = cross_attention
|
||||
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
|
||||
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
|
||||
self.decomp1 = series_decomp(moving_avg)
|
||||
self.decomp2 = series_decomp(moving_avg)
|
||||
self.decomp3 = series_decomp(moving_avg)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
|
||||
padding_mode='circular', bias=False)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
x = x + self.dropout(self.self_attention(
|
||||
x, x, x,
|
||||
attn_mask=x_mask
|
||||
)[0])
|
||||
x, trend1 = self.decomp1(x)
|
||||
x = x + self.dropout(self.cross_attention(
|
||||
x, cross, cross,
|
||||
attn_mask=cross_mask
|
||||
)[0])
|
||||
x, trend2 = self.decomp2(x)
|
||||
y = x
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
x, trend3 = self.decomp3(x + y)
|
||||
|
||||
residual_trend = trend1 + trend2 + trend3
|
||||
residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
|
||||
return x, residual_trend
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
"""
|
||||
Autoformer encoder
|
||||
"""
|
||||
def __init__(self, layers, norm_layer=None, projection=None):
|
||||
super(Decoder, self).__init__()
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.norm = norm_layer
|
||||
self.projection = projection
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
|
||||
for layer in self.layers:
|
||||
x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
|
||||
trend = trend + residual_trend
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
if self.projection is not None:
|
||||
x = self.projection(x)
|
||||
return x, trend
|
||||
-164
@@ -1,164 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils import weight_norm
|
||||
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):
|
||||
x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
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):
|
||||
x = self.value_embedding(x) + self.temporal_embedding(x_mark)
|
||||
return self.dropout(x)
|
||||
|
||||
class DataEmbedding_wo_pos_temp(nn.Module):
|
||||
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
|
||||
super(DataEmbedding_wo_pos_temp, 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):
|
||||
x = self.value_embedding(x)
|
||||
return self.dropout(x)
|
||||
|
||||
class DataEmbedding_wo_temp(nn.Module):
|
||||
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
|
||||
super(DataEmbedding_wo_temp, 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):
|
||||
x = self.value_embedding(x) + self.position_embedding(x)
|
||||
return self.dropout(x)
|
||||
@@ -1,429 +0,0 @@
|
||||
__all__ = ['PatchTST_backbone']
|
||||
|
||||
# Cell
|
||||
from typing import Callable, Optional
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
# from collections import OrderedDict
|
||||
from models.third_party.patch_tst.layers.PatchTST_layers import *
|
||||
from models.third_party.patch_tst.layers.RevIN import RevIN
|
||||
|
||||
|
||||
class CustomHead(nn.Module):
|
||||
def __init__(self, output_dim, n_vars, target_window, nf, head_dropout=0):
|
||||
super().__init__()
|
||||
self.flatten = nn.Flatten(start_dim=-3)
|
||||
self.linear = nn.Linear(nf * n_vars, output_dim * target_window)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
self.target_window = target_window
|
||||
self.output_dim = output_dim
|
||||
|
||||
def forward(self, x): # x: [bs x nvars x d_model x patch_num]
|
||||
x = self.flatten(x) # [bs x (nf * nvars)]
|
||||
x = self.linear(x) # [bs x (target_window * output_dim)]
|
||||
x = self.dropout(x)
|
||||
x = x.view(x.size(0), self.target_window, self.output_dim) # [bs x target_window x output_dim]
|
||||
# permute to match intended structure
|
||||
x = x.permute(0, 2, 1) # [bs x output_dim x target_window]
|
||||
return x
|
||||
|
||||
|
||||
# Cell
|
||||
class PatchTST_backbone(nn.Module):
|
||||
def __init__(self, c_in: int,
|
||||
context_window: int, target_window: int, patch_len: int, stride: int,
|
||||
# extras
|
||||
dec_out: int = 1,
|
||||
seq_pred: bool = False,
|
||||
#
|
||||
max_seq_len: Optional[int] = 1024,
|
||||
n_layers: int = 3, d_model=128, n_heads=16, d_k: Optional[int] = None, d_v: Optional[int] = None,
|
||||
d_ff: int = 256, norm: str = 'BatchNorm', attn_dropout: float = 0., dropout: float = 0.,
|
||||
act: str = "gelu", key_padding_mask: bool = 'auto',
|
||||
padding_var: Optional[int] = None, attn_mask: Optional[Tensor] = None, res_attention: bool = True,
|
||||
pre_norm: bool = False, store_attn: bool = False,
|
||||
pe: str = 'zeros', learn_pe: bool = True, fc_dropout: float = 0., head_dropout=0, padding_patch=None,
|
||||
pretrain_head: bool = False, head_type='flatten', individual=False, revin=True, affine=True,
|
||||
subtract_last=False,
|
||||
verbose: bool = False, **kwargs):
|
||||
|
||||
super().__init__()
|
||||
|
||||
# RevIn
|
||||
self.revin = revin
|
||||
if self.revin: self.revin_layer = RevIN(c_in, affine=affine, subtract_last=subtract_last)
|
||||
|
||||
# Patching
|
||||
self.patch_len = patch_len
|
||||
self.stride = stride
|
||||
self.padding_patch = padding_patch
|
||||
patch_num = int((context_window - patch_len) / stride + 1)
|
||||
if padding_patch == 'end': # can be modified to general case
|
||||
self.padding_patch_layer = nn.ReplicationPad1d((0, stride))
|
||||
patch_num += 1
|
||||
|
||||
# Backbone
|
||||
self.backbone = TSTiEncoder(c_in, patch_num=patch_num, patch_len=patch_len, max_seq_len=max_seq_len,
|
||||
n_layers=n_layers, d_model=d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff,
|
||||
attn_dropout=attn_dropout, dropout=dropout, act=act,
|
||||
key_padding_mask=key_padding_mask, padding_var=padding_var,
|
||||
attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
|
||||
store_attn=store_attn,
|
||||
pe=pe, learn_pe=learn_pe, verbose=verbose, **kwargs)
|
||||
|
||||
# Head
|
||||
self.head_nf = d_model * patch_num
|
||||
self.n_vars = c_in
|
||||
self.pretrain_head = pretrain_head
|
||||
self.head_type = head_type
|
||||
self.individual = individual
|
||||
# extras for non-sequence prediction
|
||||
self.seq_pred = seq_pred
|
||||
self.dec_out = dec_out
|
||||
|
||||
if self.pretrain_head:
|
||||
self.head = self.create_pretrain_head(self.head_nf, c_in,
|
||||
fc_dropout) # custom head passed as a partial func with all its kwargs
|
||||
elif not self.seq_pred:
|
||||
self.head = CustomHead(output_dim=self.dec_out,
|
||||
n_vars=self.n_vars,
|
||||
target_window=target_window,
|
||||
nf=self.head_nf,
|
||||
head_dropout=head_dropout)
|
||||
elif head_type == 'flatten':
|
||||
self.head = Flatten_Head(self.individual, self.n_vars, self.head_nf, target_window,
|
||||
head_dropout=head_dropout)
|
||||
|
||||
def forward(self, z): # z: [bs x nvars x seq_len]
|
||||
# norm
|
||||
if self.revin:
|
||||
z = z.permute(0, 2, 1)
|
||||
z = self.revin_layer(z, 'norm')
|
||||
z = z.permute(0, 2, 1)
|
||||
|
||||
# do patching
|
||||
if self.padding_patch == 'end':
|
||||
z = self.padding_patch_layer(z)
|
||||
z = z.unfold(dimension=-1, size=self.patch_len, step=self.stride) # z: [bs x nvars x patch_num x patch_len]
|
||||
z = z.permute(0, 1, 3, 2) # z: [bs x nvars x patch_len x patch_num]
|
||||
|
||||
# model
|
||||
z = self.backbone(z) # z: [bs x nvars x d_model x patch_num]
|
||||
z = self.head(z) # z: [bs x nvars x target_window]
|
||||
|
||||
# denorm
|
||||
if self.revin:
|
||||
z = z.permute(0, 2, 1)
|
||||
z = self.revin_layer(z, 'denorm')
|
||||
z = z.permute(0, 2, 1)
|
||||
|
||||
return z
|
||||
|
||||
def create_pretrain_head(self, head_nf, vars, dropout):
|
||||
return nn.Sequential(nn.Dropout(dropout),
|
||||
nn.Conv1d(head_nf, vars, 1)
|
||||
)
|
||||
|
||||
|
||||
class Flatten_Head(nn.Module):
|
||||
def __init__(self, individual, n_vars, nf, target_window, head_dropout=0):
|
||||
super().__init__()
|
||||
|
||||
self.individual = individual
|
||||
self.n_vars = n_vars
|
||||
|
||||
if self.individual:
|
||||
self.linears = nn.ModuleList()
|
||||
self.dropouts = nn.ModuleList()
|
||||
self.flattens = nn.ModuleList()
|
||||
for i in range(self.n_vars):
|
||||
self.flattens.append(nn.Flatten(start_dim=-2))
|
||||
self.linears.append(nn.Linear(nf, target_window))
|
||||
self.dropouts.append(nn.Dropout(head_dropout))
|
||||
else:
|
||||
self.flatten = nn.Flatten(start_dim=-2)
|
||||
self.linear = nn.Linear(nf, target_window)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
|
||||
def forward(self, x): # x: [bs x nvars x d_model x patch_num]
|
||||
if self.individual:
|
||||
x_out = []
|
||||
for i in range(self.n_vars):
|
||||
z = self.flattens[i](x[:, i, :, :]) # z: [bs x d_model * patch_num]
|
||||
z = self.linears[i](z) # z: [bs x target_window]
|
||||
z = self.dropouts[i](z)
|
||||
x_out.append(z)
|
||||
x = torch.stack(x_out, dim=1) # x: [bs x nvars x target_window]
|
||||
else:
|
||||
x = self.flatten(x)
|
||||
x = self.linear(x)
|
||||
x = self.dropout(x)
|
||||
return x
|
||||
|
||||
|
||||
class TSTiEncoder(nn.Module): # i means channel-independent
|
||||
def __init__(self, c_in, patch_num, patch_len, max_seq_len=1024,
|
||||
n_layers=3, d_model=128, n_heads=16, d_k=None, d_v=None,
|
||||
d_ff=256, norm='BatchNorm', attn_dropout=0., dropout=0., act="gelu", store_attn=False,
|
||||
key_padding_mask='auto', padding_var=None, attn_mask=None, res_attention=True, pre_norm=False,
|
||||
pe='zeros', learn_pe=True, verbose=False, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
self.patch_num = patch_num
|
||||
self.patch_len = patch_len
|
||||
|
||||
# Input encoding
|
||||
q_len = patch_num
|
||||
self.W_P = nn.Linear(patch_len, d_model) # Eq 1: projection of feature vectors onto a d-dim vector space
|
||||
self.seq_len = q_len
|
||||
|
||||
# Positional encoding
|
||||
self.W_pos = positional_encoding(pe, learn_pe, q_len, d_model)
|
||||
|
||||
# Residual dropout
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
# Encoder
|
||||
self.encoder = TSTEncoder(q_len, d_model, n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout, dropout=dropout,
|
||||
pre_norm=pre_norm, activation=act, res_attention=res_attention, n_layers=n_layers,
|
||||
store_attn=store_attn)
|
||||
|
||||
def forward(self, x) -> Tensor: # x: [bs x nvars x patch_len x patch_num]
|
||||
|
||||
n_vars = x.shape[1]
|
||||
# Input encoding
|
||||
x = x.permute(0, 1, 3, 2) # x: [bs x nvars x patch_num x patch_len]
|
||||
x = self.W_P(x) # x: [bs x nvars x patch_num x d_model]
|
||||
|
||||
u = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3])) # u: [bs * nvars x patch_num x d_model]
|
||||
u = self.dropout(u + self.W_pos) # u: [bs * nvars x patch_num x d_model]
|
||||
|
||||
# Encoder
|
||||
z = self.encoder(u) # z: [bs * nvars x patch_num x d_model]
|
||||
z = torch.reshape(z, (-1, n_vars, z.shape[-2], z.shape[-1])) # z: [bs x nvars x patch_num x d_model]
|
||||
z = z.permute(0, 1, 3, 2) # z: [bs x nvars x d_model x patch_num]
|
||||
|
||||
return z
|
||||
|
||||
# Cell
|
||||
|
||||
|
||||
class TSTEncoder(nn.Module):
|
||||
def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None,
|
||||
norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu',
|
||||
res_attention=False, n_layers=1, pre_norm=False, store_attn=False):
|
||||
super().__init__()
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[TSTEncoderLayer(q_len, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout, dropout=dropout,
|
||||
activation=activation, res_attention=res_attention,
|
||||
pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)])
|
||||
self.res_attention = res_attention
|
||||
|
||||
def forward(self, src: Tensor, key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
|
||||
output = src
|
||||
scores = None
|
||||
if self.res_attention:
|
||||
for mod in self.layers: output, scores = mod(output, prev=scores, key_padding_mask=key_padding_mask,
|
||||
attn_mask=attn_mask)
|
||||
return output
|
||||
else:
|
||||
for mod in self.layers: output = mod(output, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
return output
|
||||
|
||||
|
||||
class TSTEncoderLayer(nn.Module):
|
||||
def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=256, store_attn=False,
|
||||
norm='BatchNorm', attn_dropout=0, dropout=0., bias=True, activation="gelu", res_attention=False,
|
||||
pre_norm=False):
|
||||
super().__init__()
|
||||
assert not d_model % n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
|
||||
d_k = d_model // n_heads if d_k is None else d_k
|
||||
d_v = d_model // n_heads if d_v is None else d_v
|
||||
|
||||
# Multi-Head attention
|
||||
self.res_attention = res_attention
|
||||
self.self_attn = _MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout,
|
||||
proj_dropout=dropout, res_attention=res_attention)
|
||||
|
||||
# Add & Norm
|
||||
self.dropout_attn = nn.Dropout(dropout)
|
||||
if "batch" in norm.lower():
|
||||
self.norm_attn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
|
||||
else:
|
||||
self.norm_attn = nn.LayerNorm(d_model)
|
||||
|
||||
# Position-wise Feed-Forward
|
||||
self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
|
||||
get_activation_fn(activation),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(d_ff, d_model, bias=bias))
|
||||
|
||||
# Add & Norm
|
||||
self.dropout_ffn = nn.Dropout(dropout)
|
||||
if "batch" in norm.lower():
|
||||
self.norm_ffn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
|
||||
else:
|
||||
self.norm_ffn = nn.LayerNorm(d_model)
|
||||
|
||||
self.pre_norm = pre_norm
|
||||
self.store_attn = store_attn
|
||||
|
||||
def forward(self, src: Tensor, prev: Optional[Tensor] = None, key_padding_mask: Optional[Tensor] = None,
|
||||
attn_mask: Optional[Tensor] = None) -> Tensor:
|
||||
|
||||
# Multi-Head attention sublayer
|
||||
if self.pre_norm:
|
||||
src = self.norm_attn(src)
|
||||
## Multi-Head attention
|
||||
if self.res_attention:
|
||||
src2, attn, scores = self.self_attn(src, src, src, prev, key_padding_mask=key_padding_mask,
|
||||
attn_mask=attn_mask)
|
||||
else:
|
||||
src2, attn = self.self_attn(src, src, src, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
if self.store_attn:
|
||||
self.attn = attn
|
||||
## Add & Norm
|
||||
src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
|
||||
if not self.pre_norm:
|
||||
src = self.norm_attn(src)
|
||||
|
||||
# Feed-forward sublayer
|
||||
if self.pre_norm:
|
||||
src = self.norm_ffn(src)
|
||||
## Position-wise Feed-Forward
|
||||
src2 = self.ff(src)
|
||||
## Add & Norm
|
||||
src = src + self.dropout_ffn(src2) # Add: residual connection with residual dropout
|
||||
if not self.pre_norm:
|
||||
src = self.norm_ffn(src)
|
||||
|
||||
if self.res_attention:
|
||||
return src, scores
|
||||
else:
|
||||
return src
|
||||
|
||||
|
||||
class _MultiheadAttention(nn.Module):
|
||||
def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0.,
|
||||
qkv_bias=True, lsa=False):
|
||||
"""Multi Head Attention Layer
|
||||
Input shape:
|
||||
Q: [batch_size (bs) x max_q_len x d_model]
|
||||
K, V: [batch_size (bs) x q_len x d_model]
|
||||
mask: [q_len x q_len]
|
||||
"""
|
||||
super().__init__()
|
||||
d_k = d_model // n_heads if d_k is None else d_k
|
||||
d_v = d_model // n_heads if d_v is None else d_v
|
||||
|
||||
self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
|
||||
|
||||
self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
|
||||
self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
|
||||
self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)
|
||||
|
||||
# Scaled Dot-Product Attention (multiple heads)
|
||||
self.res_attention = res_attention
|
||||
self.sdp_attn = _ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout,
|
||||
res_attention=self.res_attention, lsa=lsa)
|
||||
|
||||
# Poject output
|
||||
self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))
|
||||
|
||||
def forward(self, Q: Tensor, K: Optional[Tensor] = None, V: Optional[Tensor] = None, prev: Optional[Tensor] = None,
|
||||
key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
|
||||
|
||||
bs = Q.size(0)
|
||||
if K is None: K = Q
|
||||
if V is None: V = Q
|
||||
|
||||
# Linear (+ split in multiple heads)
|
||||
q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,
|
||||
2) # q_s : [bs x n_heads x max_q_len x d_k]
|
||||
k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0, 2, 3,
|
||||
1) # k_s : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
|
||||
v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1, 2) # v_s : [bs x n_heads x q_len x d_v]
|
||||
|
||||
# Apply Scaled Dot-Product Attention (multiple heads)
|
||||
if self.res_attention:
|
||||
output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev,
|
||||
key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
else:
|
||||
output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
# output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
|
||||
|
||||
# back to the original inputs dimensions
|
||||
output = output.transpose(1, 2).contiguous().view(bs, -1,
|
||||
self.n_heads * self.d_v) # output: [bs x q_len x n_heads * d_v]
|
||||
output = self.to_out(output)
|
||||
|
||||
if self.res_attention:
|
||||
return output, attn_weights, attn_scores
|
||||
else:
|
||||
return output, attn_weights
|
||||
|
||||
|
||||
class _ScaledDotProductAttention(nn.Module):
|
||||
r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
|
||||
(Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
|
||||
by Lee et al, 2021)"""
|
||||
|
||||
def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
|
||||
super().__init__()
|
||||
self.attn_dropout = nn.Dropout(attn_dropout)
|
||||
self.res_attention = res_attention
|
||||
head_dim = d_model // n_heads
|
||||
self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
|
||||
self.lsa = lsa
|
||||
|
||||
def forward(self, q: Tensor, k: Tensor, v: Tensor, prev: Optional[Tensor] = None,
|
||||
key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
|
||||
'''
|
||||
Input shape:
|
||||
q : [bs x n_heads x max_q_len x d_k]
|
||||
k : [bs x n_heads x d_k x seq_len]
|
||||
v : [bs x n_heads x seq_len x d_v]
|
||||
prev : [bs x n_heads x q_len x seq_len]
|
||||
key_padding_mask: [bs x seq_len]
|
||||
attn_mask : [1 x seq_len x seq_len]
|
||||
Output shape:
|
||||
output: [bs x n_heads x q_len x d_v]
|
||||
attn : [bs x n_heads x q_len x seq_len]
|
||||
scores : [bs x n_heads x q_len x seq_len]
|
||||
'''
|
||||
|
||||
# Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
|
||||
attn_scores = torch.matmul(q, k) * self.scale # attn_scores : [bs x n_heads x max_q_len x q_len]
|
||||
|
||||
# Add pre-softmax attention scores from the previous layer (optional)
|
||||
if prev is not None: attn_scores = attn_scores + prev
|
||||
|
||||
# Attention mask (optional)
|
||||
if attn_mask is not None: # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_scores.masked_fill_(attn_mask, -np.inf)
|
||||
else:
|
||||
attn_scores += attn_mask
|
||||
|
||||
# Key padding mask (optional)
|
||||
if key_padding_mask is not None: # mask with shape [bs x q_len] (only when max_w_len == q_len)
|
||||
attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)
|
||||
|
||||
# normalize the attention weights
|
||||
attn_weights = F.softmax(attn_scores, dim=-1) # attn_weights : [bs x n_heads x max_q_len x q_len]
|
||||
attn_weights = self.attn_dropout(attn_weights)
|
||||
|
||||
# compute the new values given the attention weights
|
||||
output = torch.matmul(attn_weights, v) # output: [bs x n_heads x max_q_len x d_v]
|
||||
|
||||
if self.res_attention:
|
||||
return output, attn_weights, attn_scores
|
||||
else:
|
||||
return output, attn_weights
|
||||
@@ -1,121 +0,0 @@
|
||||
__all__ = ['Transpose', 'get_activation_fn', 'moving_avg', 'series_decomp', 'PositionalEncoding', 'SinCosPosEncoding', 'Coord2dPosEncoding', 'Coord1dPosEncoding', 'positional_encoding']
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import math
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self, *dims, contiguous=False):
|
||||
super().__init__()
|
||||
self.dims, self.contiguous = dims, contiguous
|
||||
def forward(self, x):
|
||||
if self.contiguous: return x.transpose(*self.dims).contiguous()
|
||||
else: return x.transpose(*self.dims)
|
||||
|
||||
|
||||
def get_activation_fn(activation):
|
||||
if callable(activation): return activation()
|
||||
elif activation.lower() == "relu": return nn.ReLU()
|
||||
elif activation.lower() == "gelu": return nn.GELU()
|
||||
raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable')
|
||||
|
||||
|
||||
# decomposition
|
||||
|
||||
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
|
||||
|
||||
|
||||
|
||||
# pos_encoding
|
||||
|
||||
def PositionalEncoding(q_len, d_model, normalize=True):
|
||||
pe = torch.zeros(q_len, d_model)
|
||||
position = torch.arange(0, q_len).unsqueeze(1)
|
||||
div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
|
||||
pe[:, 0::2] = torch.sin(position * div_term)
|
||||
pe[:, 1::2] = torch.cos(position * div_term)
|
||||
if normalize:
|
||||
pe = pe - pe.mean()
|
||||
pe = pe / (pe.std() * 10)
|
||||
return pe
|
||||
|
||||
SinCosPosEncoding = PositionalEncoding
|
||||
|
||||
def Coord2dPosEncoding(q_len, d_model, exponential=False, normalize=True, eps=1e-3, verbose=False):
|
||||
x = .5 if exponential else 1
|
||||
i = 0
|
||||
for i in range(100):
|
||||
cpe = 2 * (torch.linspace(0, 1, q_len).reshape(-1, 1) ** x) * (torch.linspace(0, 1, d_model).reshape(1, -1) ** x) - 1
|
||||
pv(f'{i:4.0f} {x:5.3f} {cpe.mean():+6.3f}', verbose)
|
||||
if abs(cpe.mean()) <= eps: break
|
||||
elif cpe.mean() > eps: x += .001
|
||||
else: x -= .001
|
||||
i += 1
|
||||
if normalize:
|
||||
cpe = cpe - cpe.mean()
|
||||
cpe = cpe / (cpe.std() * 10)
|
||||
return cpe
|
||||
|
||||
def Coord1dPosEncoding(q_len, exponential=False, normalize=True):
|
||||
cpe = (2 * (torch.linspace(0, 1, q_len).reshape(-1, 1)**(.5 if exponential else 1)) - 1)
|
||||
if normalize:
|
||||
cpe = cpe - cpe.mean()
|
||||
cpe = cpe / (cpe.std() * 10)
|
||||
return cpe
|
||||
|
||||
def positional_encoding(pe, learn_pe, q_len, d_model):
|
||||
# Positional encoding
|
||||
if pe == None:
|
||||
W_pos = torch.empty((q_len, d_model)) # pe = None and learn_pe = False can be used to measure impact of pe
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
learn_pe = False
|
||||
elif pe == 'zero':
|
||||
W_pos = torch.empty((q_len, 1))
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
elif pe == 'zeros':
|
||||
W_pos = torch.empty((q_len, d_model))
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
elif pe == 'normal' or pe == 'gauss':
|
||||
W_pos = torch.zeros((q_len, 1))
|
||||
torch.nn.init.normal_(W_pos, mean=0.0, std=0.1)
|
||||
elif pe == 'uniform':
|
||||
W_pos = torch.zeros((q_len, 1))
|
||||
nn.init.uniform_(W_pos, a=0.0, b=0.1)
|
||||
elif pe == 'lin1d': W_pos = Coord1dPosEncoding(q_len, exponential=False, normalize=True)
|
||||
elif pe == 'exp1d': W_pos = Coord1dPosEncoding(q_len, exponential=True, normalize=True)
|
||||
elif pe == 'lin2d': W_pos = Coord2dPosEncoding(q_len, d_model, exponential=False, normalize=True)
|
||||
elif pe == 'exp2d': W_pos = Coord2dPosEncoding(q_len, d_model, exponential=True, normalize=True)
|
||||
elif pe == 'sincos': W_pos = PositionalEncoding(q_len, d_model, normalize=True)
|
||||
else: raise ValueError(f"{pe} is not a valid pe (positional encoder. Available types: 'gauss'=='normal', \
|
||||
'zeros', 'zero', uniform', 'lin1d', 'exp1d', 'lin2d', 'exp2d', 'sincos', None.)")
|
||||
return nn.Parameter(W_pos, requires_grad=learn_pe)
|
||||
@@ -1,63 +0,0 @@
|
||||
# code from https://github.com/ts-kim/RevIN, with minor modifications
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
class RevIN(nn.Module):
|
||||
def __init__(self, num_features: int, eps=1e-5, affine=True, subtract_last=False):
|
||||
"""
|
||||
:param num_features: the number of features or channels
|
||||
:param eps: a value added for numerical stability
|
||||
:param affine: if True, RevIN has learnable affine parameters
|
||||
"""
|
||||
super(RevIN, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
self.subtract_last = subtract_last
|
||||
if self.affine:
|
||||
self._init_params()
|
||||
|
||||
def forward(self, x, mode:str):
|
||||
if mode == 'norm':
|
||||
self._get_statistics(x)
|
||||
x = self._normalize(x)
|
||||
elif mode == 'denorm':
|
||||
x = self._denormalize(x)
|
||||
else: raise NotImplementedError
|
||||
return x
|
||||
|
||||
def _init_params(self):
|
||||
# initialize RevIN params: (C,)
|
||||
self.affine_weight = nn.Parameter(torch.ones(self.num_features))
|
||||
self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
|
||||
|
||||
def _get_statistics(self, x):
|
||||
dim2reduce = tuple(range(1, x.ndim-1))
|
||||
if self.subtract_last:
|
||||
self.last = x[:,-1,:].unsqueeze(1)
|
||||
else:
|
||||
self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
|
||||
self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
|
||||
|
||||
def _normalize(self, x):
|
||||
if self.subtract_last:
|
||||
x = x - self.last
|
||||
else:
|
||||
x = x - self.mean
|
||||
x = x / self.stdev
|
||||
if self.affine:
|
||||
x = x * self.affine_weight
|
||||
x = x + self.affine_bias
|
||||
return x
|
||||
|
||||
def _denormalize(self, x):
|
||||
if self.affine:
|
||||
x = x - self.affine_bias
|
||||
x = x / (self.affine_weight + self.eps*self.eps)
|
||||
x = x * self.stdev
|
||||
if self.subtract_last:
|
||||
x = x + self.last
|
||||
else:
|
||||
x = x + self.mean
|
||||
return x
|
||||
@@ -1,166 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import numpy as np
|
||||
import math
|
||||
from math import sqrt
|
||||
from utils.masking import TriangularCausalMask, ProbMask
|
||||
import os
|
||||
|
||||
|
||||
class FullAttention(nn.Module):
|
||||
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(FullAttention, self).__init__()
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
scale = self.scale or 1. / sqrt(E)
|
||||
|
||||
scores = torch.einsum("blhe,bshe->bhls", queries, keys)
|
||||
|
||||
if self.mask_flag:
|
||||
if attn_mask is None:
|
||||
attn_mask = TriangularCausalMask(B, L, device=queries.device)
|
||||
|
||||
scores.masked_fill_(attn_mask.mask, -np.inf)
|
||||
|
||||
A = self.dropout(torch.softmax(scale * scores, dim=-1))
|
||||
V = torch.einsum("bhls,bshd->blhd", A, values)
|
||||
|
||||
if self.output_attention:
|
||||
return (V.contiguous(), A)
|
||||
else:
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class ProbAttention(nn.Module):
|
||||
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(ProbAttention, self).__init__()
|
||||
self.factor = factor
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def _prob_QK(self, Q, K, sample_k, n_top): # n_top: c*ln(L_q)
|
||||
# Q [B, H, L, D]
|
||||
B, H, L_K, E = K.shape
|
||||
_, _, L_Q, _ = Q.shape
|
||||
|
||||
# calculate the sampled Q_K
|
||||
K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
|
||||
index_sample = torch.randint(L_K, (L_Q, sample_k)) # real U = U_part(factor*ln(L_k))*L_q
|
||||
K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]
|
||||
Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
|
||||
|
||||
# find the Top_k query with sparisty measurement
|
||||
M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
|
||||
M_top = M.topk(n_top, sorted=False)[1]
|
||||
|
||||
# use the reduced Q to calculate Q_K
|
||||
Q_reduce = Q[torch.arange(B)[:, None, None],
|
||||
torch.arange(H)[None, :, None],
|
||||
M_top, :] # factor*ln(L_q)
|
||||
Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) # factor*ln(L_q)*L_k
|
||||
|
||||
return Q_K, M_top
|
||||
|
||||
def _get_initial_context(self, V, L_Q):
|
||||
B, H, L_V, D = V.shape
|
||||
if not self.mask_flag:
|
||||
# V_sum = V.sum(dim=-2)
|
||||
V_sum = V.mean(dim=-2)
|
||||
contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
|
||||
else: # use mask
|
||||
assert (L_Q == L_V) # requires that L_Q == L_V, i.e. for self-attention only
|
||||
contex = V.cumsum(dim=-2)
|
||||
return contex
|
||||
|
||||
def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
|
||||
B, H, L_V, D = V.shape
|
||||
|
||||
if self.mask_flag:
|
||||
attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
|
||||
scores.masked_fill_(attn_mask.mask, -np.inf)
|
||||
|
||||
attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores)
|
||||
|
||||
context_in[torch.arange(B)[:, None, None],
|
||||
torch.arange(H)[None, :, None],
|
||||
index, :] = torch.matmul(attn, V).type_as(context_in)
|
||||
if self.output_attention:
|
||||
attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)
|
||||
attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn
|
||||
return (context_in, attns)
|
||||
else:
|
||||
return (context_in, None)
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L_Q, H, D = queries.shape
|
||||
_, L_K, _, _ = keys.shape
|
||||
|
||||
queries = queries.transpose(2, 1)
|
||||
keys = keys.transpose(2, 1)
|
||||
values = values.transpose(2, 1)
|
||||
|
||||
U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item() # c*ln(L_k)
|
||||
u = self.factor * np.ceil(np.log(L_Q)).astype('int').item() # c*ln(L_q)
|
||||
|
||||
U_part = U_part if U_part < L_K else L_K
|
||||
u = u if u < L_Q else L_Q
|
||||
|
||||
scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
|
||||
|
||||
# add scale factor
|
||||
scale = self.scale or 1. / sqrt(D)
|
||||
if scale is not None:
|
||||
scores_top = scores_top * scale
|
||||
# get the context
|
||||
context = self._get_initial_context(values, L_Q)
|
||||
# update the context with selected top_k queries
|
||||
context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
|
||||
|
||||
return context.contiguous(), attn
|
||||
|
||||
|
||||
class AttentionLayer(nn.Module):
|
||||
def __init__(self, attention, d_model, n_heads, d_keys=None,
|
||||
d_values=None):
|
||||
super(AttentionLayer, self).__init__()
|
||||
|
||||
d_keys = d_keys or (d_model // n_heads)
|
||||
d_values = d_values or (d_model // n_heads)
|
||||
|
||||
self.inner_attention = attention
|
||||
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.value_projection = nn.Linear(d_model, d_values * n_heads)
|
||||
self.out_projection = nn.Linear(d_values * n_heads, d_model)
|
||||
self.n_heads = n_heads
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, _ = queries.shape
|
||||
_, S, _ = keys.shape
|
||||
H = self.n_heads
|
||||
|
||||
queries = self.query_projection(queries).view(B, L, H, -1)
|
||||
keys = self.key_projection(keys).view(B, S, H, -1)
|
||||
values = self.value_projection(values).view(B, S, H, -1)
|
||||
|
||||
out, attn = self.inner_attention(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
attn_mask
|
||||
)
|
||||
out = out.view(B, L, -1)
|
||||
|
||||
return self.out_projection(out), attn
|
||||
@@ -1,131 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(self, c_in):
|
||||
super(ConvLayer, self).__init__()
|
||||
self.downConv = nn.Conv1d(in_channels=c_in,
|
||||
out_channels=c_in,
|
||||
kernel_size=3,
|
||||
padding=2,
|
||||
padding_mode='circular')
|
||||
self.norm = nn.BatchNorm1d(c_in)
|
||||
self.activation = nn.ELU()
|
||||
self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.downConv(x.permute(0, 2, 1))
|
||||
x = self.norm(x)
|
||||
x = self.activation(x)
|
||||
x = self.maxPool(x)
|
||||
x = x.transpose(1, 2)
|
||||
return x
|
||||
|
||||
|
||||
class EncoderLayer(nn.Module):
|
||||
def __init__(self, attention, 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.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
|
||||
|
||||
def forward(self, x, attn_mask=None):
|
||||
new_x, attn = self.attention(
|
||||
x, x, x,
|
||||
attn_mask=attn_mask
|
||||
)
|
||||
x = x + self.dropout(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):
|
||||
# x [B, L, D]
|
||||
attns = []
|
||||
if self.conv_layers is not None:
|
||||
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
x = conv_layer(x)
|
||||
attns.append(attn)
|
||||
x, attn = self.attn_layers[-1](x)
|
||||
attns.append(attn)
|
||||
else:
|
||||
for attn_layer in self.attn_layers:
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
attns.append(attn)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
return x, attns
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
|
||||
dropout=0.1, activation="relu"):
|
||||
super(DecoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.self_attention = self_attention
|
||||
self.cross_attention = cross_attention
|
||||
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.norm3 = nn.LayerNorm(d_model)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
x = x + self.dropout(self.self_attention(
|
||||
x, x, x,
|
||||
attn_mask=x_mask
|
||||
)[0])
|
||||
x = self.norm1(x)
|
||||
|
||||
x = x + self.dropout(self.cross_attention(
|
||||
x, cross, cross,
|
||||
attn_mask=cross_mask
|
||||
)[0])
|
||||
|
||||
y = x = self.norm2(x)
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
|
||||
return self.norm3(x + y)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, layers, norm_layer=None, projection=None):
|
||||
super(Decoder, self).__init__()
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.norm = norm_layer
|
||||
self.projection = projection
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
for layer in self.layers:
|
||||
x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
if self.projection is not None:
|
||||
x = self.projection(x)
|
||||
return x
|
||||
@@ -1,121 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from layers.Embed import DataEmbedding, DataEmbedding_wo_pos,DataEmbedding_wo_pos_temp,DataEmbedding_wo_temp
|
||||
from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
|
||||
from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Autoformer is the first method to achieve the series-wise connection,
|
||||
with inherent O(LlogL) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.label_len = configs.label_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Decomp
|
||||
kernel_size = configs.moving_avg
|
||||
self.decomp = series_decomp(kernel_size)
|
||||
|
||||
# Embedding
|
||||
# The series-wise connection inherently contains the sequential information.
|
||||
# Thus, we can discard the position embedding of transformers.
|
||||
if configs.embed_type == 0:
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 1:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 2:
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
elif configs.embed_type == 3:
|
||||
self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 4:
|
||||
self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(True, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.c_out,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
# decomp init
|
||||
mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
|
||||
zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)
|
||||
seasonal_init, trend_init = self.decomp(x_enc)
|
||||
# decoder input
|
||||
trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
|
||||
seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)
|
||||
# enc
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
# dec
|
||||
dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
|
||||
seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
|
||||
trend=trend_init)
|
||||
# final
|
||||
dec_out = trend_part + seasonal_part
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
@@ -1,87 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
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 Model(nn.Module):
|
||||
"""
|
||||
Decomposition-Linear
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.pred_len = configs.pred_len
|
||||
|
||||
# Decompsition Kernel Size
|
||||
kernel_size = 25
|
||||
self.decompsition = series_decomp(kernel_size)
|
||||
self.individual = configs.individual
|
||||
self.channels = configs.enc_in
|
||||
|
||||
if self.individual:
|
||||
self.Linear_Seasonal = nn.ModuleList()
|
||||
self.Linear_Trend = nn.ModuleList()
|
||||
|
||||
for i in range(self.channels):
|
||||
self.Linear_Seasonal.append(nn.Linear(self.seq_len,self.pred_len))
|
||||
self.Linear_Trend.append(nn.Linear(self.seq_len,self.pred_len))
|
||||
|
||||
# Use this two lines if you want to visualize the weights
|
||||
# self.Linear_Seasonal[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
# self.Linear_Trend[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
else:
|
||||
self.Linear_Seasonal = nn.Linear(self.seq_len,self.pred_len)
|
||||
self.Linear_Trend = nn.Linear(self.seq_len,self.pred_len)
|
||||
|
||||
# Use this two lines if you want to visualize the weights
|
||||
# self.Linear_Seasonal.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
# self.Linear_Trend.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
|
||||
def forward(self, x):
|
||||
# x: [Batch, Input length, Channel]
|
||||
seasonal_init, trend_init = self.decompsition(x)
|
||||
seasonal_init, trend_init = seasonal_init.permute(0,2,1), trend_init.permute(0,2,1)
|
||||
if self.individual:
|
||||
seasonal_output = torch.zeros([seasonal_init.size(0),seasonal_init.size(1),self.pred_len],dtype=seasonal_init.dtype).to(seasonal_init.device)
|
||||
trend_output = torch.zeros([trend_init.size(0),trend_init.size(1),self.pred_len],dtype=trend_init.dtype).to(trend_init.device)
|
||||
for i in range(self.channels):
|
||||
seasonal_output[:,i,:] = self.Linear_Seasonal[i](seasonal_init[:,i,:])
|
||||
trend_output[:,i,:] = self.Linear_Trend[i](trend_init[:,i,:])
|
||||
else:
|
||||
seasonal_output = self.Linear_Seasonal(seasonal_init)
|
||||
trend_output = self.Linear_Trend(trend_init)
|
||||
|
||||
x = seasonal_output + trend_output
|
||||
return x.permute(0,2,1) # to [Batch, Output length, Channel]
|
||||
@@ -1,101 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from utils.masking import TriangularCausalMask, ProbMask
|
||||
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
|
||||
from layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer
|
||||
from layers.Embed import DataEmbedding,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Informer with Propspare attention in O(LlogL) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Embedding
|
||||
if configs.embed_type == 0:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 1:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 2:
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
elif configs.embed_type == 3:
|
||||
self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 4:
|
||||
self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AttentionLayer(
|
||||
ProbAttention(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
[
|
||||
ConvLayer(
|
||||
configs.d_model
|
||||
) for l in range(configs.e_layers - 1)
|
||||
] if configs.distil else None,
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AttentionLayer(
|
||||
ProbAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AttentionLayer(
|
||||
ProbAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
|
||||
dec_out = self.dec_embedding(x_dec, x_mark_dec)
|
||||
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
@@ -1,21 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Just one Linear layer
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.Linear = nn.Linear(self.seq_len, self.pred_len)
|
||||
# Use this line if you want to visualize the weights
|
||||
# self.Linear.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
|
||||
def forward(self, x):
|
||||
# x: [Batch, Input length, Channel]
|
||||
x = self.Linear(x.permute(0,2,1)).permute(0,2,1)
|
||||
return x # [Batch, Output length, Channel]
|
||||
@@ -1,24 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Normalization-Linear
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.Linear = nn.Linear(self.seq_len, self.pred_len)
|
||||
# Use this line if you want to visualize the weights
|
||||
# self.Linear.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
|
||||
|
||||
def forward(self, x):
|
||||
# x: [Batch, Input length, Channel]
|
||||
seq_last = x[:,-1:,:].detach()
|
||||
x = x - seq_last
|
||||
x = self.Linear(x.permute(0,2,1)).permute(0,2,1)
|
||||
x = x + seq_last
|
||||
return x # [Batch, Output length, Channel]
|
||||
@@ -1,127 +0,0 @@
|
||||
__all__ = ['PatchTST']
|
||||
|
||||
# Cell
|
||||
from typing import Callable, Optional
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
from models.third_party.patch_tst.layers.PatchTST_backbone import PatchTST_backbone
|
||||
from models.third_party.patch_tst.layers.PatchTST_layers import series_decomp
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
def __init__(self, input_dim: int, output_dim: int, configs, max_seq_len: Optional[int] = 1024,
|
||||
d_k: Optional[int] = None, d_v: Optional[int] = None,
|
||||
norm: str = 'BatchNorm', attn_dropout: float = 0.,
|
||||
act: str = "gelu", key_padding_mask: bool = 'auto', padding_var: Optional[int] = None,
|
||||
attn_mask: Optional[Tensor] = None, res_attention: bool = True,
|
||||
pre_norm: bool = False, store_attn: bool = False, pe: str = 'zeros', learn_pe: bool = True,
|
||||
pretrain_head: bool = False, head_type='flatten', verbose: bool = False, **kwargs):
|
||||
|
||||
super().__init__()
|
||||
|
||||
# load parameters
|
||||
c_in = input_dim
|
||||
context_window = configs['seq_len']
|
||||
target_window = configs['pred_len']
|
||||
dec_out = output_dim
|
||||
seq_pred = configs["seq_pred"]
|
||||
|
||||
n_layers = configs['e_layers']
|
||||
n_heads = configs['n_heads']
|
||||
d_model = configs['d_model']
|
||||
d_ff = configs['d_ff']
|
||||
dropout = configs['dropout']
|
||||
fc_dropout = configs['fc_dropout']
|
||||
head_dropout = configs['head_dropout']
|
||||
|
||||
individual = configs['individual']
|
||||
|
||||
patch_len = configs['patch_len']
|
||||
stride = configs['stride']
|
||||
padding_patch = configs['padding_patch']
|
||||
|
||||
revin = configs['revin']
|
||||
affine = configs['affine']
|
||||
subtract_last = configs['subtract_last']
|
||||
|
||||
decomposition = configs['decomposition']
|
||||
kernel_size = configs['kernel_size']
|
||||
|
||||
# model
|
||||
self.decomposition = decomposition
|
||||
if self.decomposition:
|
||||
self.decomp_module = series_decomp(kernel_size)
|
||||
self.model_trend = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
|
||||
# extras
|
||||
dec_out=dec_out,
|
||||
seq_pred=seq_pred,
|
||||
#
|
||||
patch_len=patch_len, stride=stride,
|
||||
max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
|
||||
n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout,
|
||||
dropout=dropout, act=act, key_padding_mask=key_padding_mask,
|
||||
padding_var=padding_var,
|
||||
attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
|
||||
store_attn=store_attn,
|
||||
pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
|
||||
head_dropout=head_dropout, padding_patch=padding_patch,
|
||||
pretrain_head=pretrain_head, head_type=head_type,
|
||||
individual=individual, revin=revin, affine=affine,
|
||||
subtract_last=subtract_last, verbose=verbose, **kwargs)
|
||||
self.model_res = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
|
||||
# extras
|
||||
dec_out=dec_out,
|
||||
seq_pred=seq_pred,
|
||||
#
|
||||
patch_len=patch_len, stride=stride,
|
||||
max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
|
||||
n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout,
|
||||
dropout=dropout, act=act, key_padding_mask=key_padding_mask,
|
||||
padding_var=padding_var,
|
||||
attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
|
||||
store_attn=store_attn,
|
||||
pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
|
||||
head_dropout=head_dropout, padding_patch=padding_patch,
|
||||
pretrain_head=pretrain_head, head_type=head_type, individual=individual,
|
||||
revin=revin, affine=affine,
|
||||
subtract_last=subtract_last, verbose=verbose, **kwargs)
|
||||
else:
|
||||
self.model = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
|
||||
# extras
|
||||
dec_out=dec_out,
|
||||
seq_pred=seq_pred,
|
||||
#
|
||||
patch_len=patch_len, stride=stride,
|
||||
max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
|
||||
n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout,
|
||||
dropout=dropout, act=act, key_padding_mask=key_padding_mask,
|
||||
padding_var=padding_var,
|
||||
attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
|
||||
store_attn=store_attn,
|
||||
pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout, head_dropout=head_dropout,
|
||||
padding_patch=padding_patch,
|
||||
pretrain_head=pretrain_head, head_type=head_type, individual=individual,
|
||||
revin=revin, affine=affine,
|
||||
subtract_last=subtract_last, verbose=verbose, **kwargs)
|
||||
|
||||
def forward(self, x): # x: [Batch, Input length, Channel]
|
||||
if self.decomposition:
|
||||
res_init, trend_init = self.decomp_module(x)
|
||||
res_init, trend_init = res_init.permute(0, 2, 1), trend_init.permute(0, 2,
|
||||
1) # x: [Batch, Channel, Input length]
|
||||
res = self.model_res(res_init)
|
||||
trend = self.model_trend(trend_init)
|
||||
x = res + trend
|
||||
x = x.permute(0, 2, 1) # x: [Batch, Input length, Channel]
|
||||
else:
|
||||
x = x.permute(0, 2, 1) # x: [Batch, Channel, Input length]
|
||||
x = self.model(x)
|
||||
x = x.permute(0, 2, 1) # x: [Batch, Input length, Channel]
|
||||
return x
|
||||
@@ -1,120 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
import pmdarima as pm
|
||||
import threading
|
||||
from sklearn.ensemble import GradientBoostingRegressor
|
||||
|
||||
class Naive_repeat(nn.Module):
|
||||
def __init__(self, configs):
|
||||
super(Naive_repeat, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
|
||||
def forward(self, x):
|
||||
B,L,D = x.shape
|
||||
x = x[:,-1,:].reshape(B,1,D).repeat(self.pred_len,axis=1)
|
||||
return x # [B, L, D]
|
||||
|
||||
class Naive_thread(threading.Thread):
|
||||
def __init__(self,func,args=()):
|
||||
super(Naive_thread,self).__init__()
|
||||
self.func = func
|
||||
self.args = args
|
||||
|
||||
def run(self):
|
||||
self.results = self.func(*self.args)
|
||||
|
||||
def return_result(self):
|
||||
threading.Thread.join(self)
|
||||
return self.results
|
||||
|
||||
def _arima(seq,pred_len,bt,i):
|
||||
model = pm.auto_arima(seq)
|
||||
forecasts = model.predict(pred_len)
|
||||
return forecasts,bt,i
|
||||
|
||||
class Arima(nn.Module):
|
||||
"""
|
||||
Extremely slow, please sample < 0.1
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Arima, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
|
||||
def forward(self, x):
|
||||
result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
|
||||
threads = []
|
||||
for bt,seqs in tqdm(enumerate(x)):
|
||||
for i in range(seqs.shape[-1]):
|
||||
seq = seqs[:,i]
|
||||
one_seq = Naive_thread(func=_arima,args=(seq,self.pred_len,bt,i))
|
||||
threads.append(one_seq)
|
||||
threads[-1].start()
|
||||
for every_thread in tqdm(threads):
|
||||
forcast,bt,i = every_thread.return_result()
|
||||
result[bt,:,i] = forcast
|
||||
|
||||
return result # [B, L, D]
|
||||
|
||||
def _sarima(season,seq,pred_len,bt,i):
|
||||
model = pm.auto_arima(seq, seasonal=True, m=season)
|
||||
forecasts = model.predict(pred_len)
|
||||
return forecasts,bt,i
|
||||
|
||||
class SArima(nn.Module):
|
||||
"""
|
||||
Extremely extremely slow, please sample < 0.01
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(SArima, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
self.seq_len = configs.seq_len
|
||||
self.season = 24
|
||||
if 'Ettm' in configs.data_path:
|
||||
self.season = 12
|
||||
elif 'ILI' in configs.data_path:
|
||||
self.season = 1
|
||||
if self.season >= self.seq_len:
|
||||
self.season = 1
|
||||
|
||||
def forward(self, x):
|
||||
result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
|
||||
threads = []
|
||||
for bt,seqs in tqdm(enumerate(x)):
|
||||
for i in range(seqs.shape[-1]):
|
||||
seq = seqs[:,i]
|
||||
one_seq = Naive_thread(func=_sarima,args=(self.season,seq,self.pred_len,bt,i))
|
||||
threads.append(one_seq)
|
||||
threads[-1].start()
|
||||
for every_thread in tqdm(threads):
|
||||
forcast,bt,i = every_thread.return_result()
|
||||
result[bt,:,i] = forcast
|
||||
return result # [B, L, D]
|
||||
|
||||
def _gbrt(seq,seq_len,pred_len,bt,i):
|
||||
model = GradientBoostingRegressor()
|
||||
model.fit(np.arange(seq_len).reshape(-1,1),seq.reshape(-1,1))
|
||||
forecasts = model.predict(np.arange(seq_len,seq_len+pred_len).reshape(-1,1))
|
||||
return forecasts,bt,i
|
||||
|
||||
class GBRT(nn.Module):
|
||||
def __init__(self, configs):
|
||||
super(GBRT, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.pred_len = configs.pred_len
|
||||
|
||||
def forward(self, x):
|
||||
result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
|
||||
threads = []
|
||||
for bt,seqs in tqdm(enumerate(x)):
|
||||
for i in range(seqs.shape[-1]):
|
||||
seq = seqs[:,i]
|
||||
one_seq = Naive_thread(func=_gbrt,args=(seq,self.seq_len,self.pred_len,bt,i))
|
||||
threads.append(one_seq)
|
||||
threads[-1].start()
|
||||
for every_thread in tqdm(threads):
|
||||
forcast,bt,i = every_thread.return_result()
|
||||
result[bt,:,i] = forcast
|
||||
return result # [B, L, D]
|
||||
@@ -1,94 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
|
||||
from layers.SelfAttention_Family import FullAttention, AttentionLayer
|
||||
from layers.Embed import DataEmbedding,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Vanilla Transformer with O(L^2) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Embedding
|
||||
if configs.embed_type == 0:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 1:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 2:
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
elif configs.embed_type == 3:
|
||||
self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 4:
|
||||
self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AttentionLayer(
|
||||
FullAttention(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention), configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AttentionLayer(
|
||||
FullAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AttentionLayer(
|
||||
FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
|
||||
dec_out = self.dec_embedding(x_dec, x_mark_dec)
|
||||
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
@@ -1,2 +0,0 @@
|
||||
* text=auto
|
||||
*.sh text eol=lf
|
||||
@@ -1,465 +0,0 @@
|
||||
# Created by https://www.toptal.com/developers/gitignore/api/emacs,jetbrains+all,jupyternotebooks,linux,macos,microsoftoffice,notepadpp,python,tortoisegit,vim,visualstudiocode,windows
|
||||
# Edit at https://www.toptal.com/developers/gitignore?templates=emacs,jetbrains+all,jupyternotebooks,linux,macos,microsoftoffice,notepadpp,python,tortoisegit,vim,visualstudiocode,windows
|
||||
|
||||
### Emacs ###
|
||||
# -*- mode: gitignore; -*-
|
||||
*~
|
||||
\#*\#
|
||||
/.emacs.desktop
|
||||
/.emacs.desktop.lock
|
||||
*.elc
|
||||
auto-save-list
|
||||
tramp
|
||||
.\#*
|
||||
|
||||
# Org-mode
|
||||
.org-id-locations
|
||||
*_archive
|
||||
|
||||
# flymake-mode
|
||||
*_flymake.*
|
||||
|
||||
# eshell files
|
||||
/eshell/history
|
||||
/eshell/lastdir
|
||||
|
||||
# elpa packages
|
||||
/elpa/
|
||||
|
||||
# reftex files
|
||||
*.rel
|
||||
|
||||
# AUCTeX auto folder
|
||||
/auto/
|
||||
|
||||
# cask packages
|
||||
.cask/
|
||||
dist/
|
||||
|
||||
# Flycheck
|
||||
flycheck_*.el
|
||||
|
||||
# server auth directory
|
||||
/server/
|
||||
|
||||
# projectiles files
|
||||
.projectile
|
||||
|
||||
# directory configuration
|
||||
.dir-locals.el
|
||||
|
||||
# network security
|
||||
/network-security.data
|
||||
|
||||
|
||||
### JetBrains+all ###
|
||||
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
|
||||
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
|
||||
|
||||
# User-specific stuff
|
||||
.idea/**/workspace.xml
|
||||
.idea/**/tasks.xml
|
||||
.idea/**/usage.statistics.xml
|
||||
.idea/**/dictionaries
|
||||
.idea/**/shelf
|
||||
|
||||
# AWS User-specific
|
||||
.idea/**/aws.xml
|
||||
|
||||
# Generated files
|
||||
.idea/**/contentModel.xml
|
||||
|
||||
# Sensitive or high-churn files
|
||||
.idea/**/dataSources/
|
||||
.idea/**/dataSources.ids
|
||||
.idea/**/dataSources.local.xml
|
||||
.idea/**/sqlDataSources.xml
|
||||
.idea/**/dynamic.xml
|
||||
.idea/**/uiDesigner.xml
|
||||
.idea/**/dbnavigator.xml
|
||||
|
||||
# Gradle
|
||||
.idea/**/gradle.xml
|
||||
.idea/**/libraries
|
||||
|
||||
# Gradle and Maven with auto-import
|
||||
# When using Gradle or Maven with auto-import, you should exclude module files,
|
||||
# since they will be recreated, and may cause churn. Uncomment if using
|
||||
# auto-import.
|
||||
# .idea/artifacts
|
||||
# .idea/compiler.xml
|
||||
# .idea/jarRepositories.xml
|
||||
# .idea/modules.xml
|
||||
# .idea/*.iml
|
||||
# .idea/modules
|
||||
# *.iml
|
||||
# *.ipr
|
||||
|
||||
# CMake
|
||||
cmake-build-*/
|
||||
|
||||
# Mongo Explorer plugin
|
||||
.idea/**/mongoSettings.xml
|
||||
|
||||
# File-based project format
|
||||
*.iws
|
||||
|
||||
# IntelliJ
|
||||
out/
|
||||
|
||||
# mpeltonen/sbt-idea plugin
|
||||
.idea_modules/
|
||||
|
||||
# JIRA plugin
|
||||
atlassian-ide-plugin.xml
|
||||
|
||||
# Cursive Clojure plugin
|
||||
.idea/replstate.xml
|
||||
|
||||
# SonarLint plugin
|
||||
.idea/sonarlint/
|
||||
|
||||
# Crashlytics plugin (for Android Studio and IntelliJ)
|
||||
com_crashlytics_export_strings.xml
|
||||
crashlytics.properties
|
||||
crashlytics-build.properties
|
||||
fabric.properties
|
||||
|
||||
# Editor-based Rest Client
|
||||
.idea/httpRequests
|
||||
|
||||
# Android studio 3.1+ serialized cache file
|
||||
.idea/caches/build_file_checksums.ser
|
||||
|
||||
### JetBrains+all Patch ###
|
||||
# Ignore everything but code style settings and run configurations
|
||||
# that are supposed to be shared within teams.
|
||||
|
||||
.idea/*
|
||||
|
||||
!.idea/codeStyles
|
||||
!.idea/runConfigurations
|
||||
|
||||
### JupyterNotebooks ###
|
||||
# gitignore template for Jupyter Notebooks
|
||||
# website: http://jupyter.org/
|
||||
|
||||
.ipynb_checkpoints
|
||||
*/.ipynb_checkpoints/*
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# Remove previous ipynb_checkpoints
|
||||
# git rm -r .ipynb_checkpoints/
|
||||
|
||||
### Linux ###
|
||||
|
||||
# temporary files which can be created if a process still has a handle open of a deleted file
|
||||
.fuse_hidden*
|
||||
|
||||
# KDE directory preferences
|
||||
.directory
|
||||
|
||||
# Linux trash folder which might appear on any partition or disk
|
||||
.Trash-*
|
||||
|
||||
# .nfs files are created when an open file is removed but is still being accessed
|
||||
.nfs*
|
||||
|
||||
### macOS ###
|
||||
# General
|
||||
.DS_Store
|
||||
.AppleDouble
|
||||
.LSOverride
|
||||
|
||||
# Icon must end with two \r
|
||||
Icon
|
||||
|
||||
|
||||
# Thumbnails
|
||||
._*
|
||||
|
||||
# Files that might appear in the root of a volume
|
||||
.DocumentRevisions-V100
|
||||
.fseventsd
|
||||
.Spotlight-V100
|
||||
.TemporaryItems
|
||||
.Trashes
|
||||
.VolumeIcon.icns
|
||||
.com.apple.timemachine.donotpresent
|
||||
|
||||
# Directories potentially created on remote AFP share
|
||||
.AppleDB
|
||||
.AppleDesktop
|
||||
Network Trash Folder
|
||||
Temporary Items
|
||||
.apdisk
|
||||
|
||||
### macOS Patch ###
|
||||
# iCloud generated files
|
||||
*.icloud
|
||||
|
||||
### MicrosoftOffice ###
|
||||
*.tmp
|
||||
|
||||
# Word temporary
|
||||
~$*.doc*
|
||||
|
||||
# Word Auto Backup File
|
||||
Backup of *.doc*
|
||||
|
||||
# Excel temporary
|
||||
~$*.xls*
|
||||
|
||||
# Excel Backup File
|
||||
*.xlk
|
||||
|
||||
# PowerPoint temporary
|
||||
~$*.ppt*
|
||||
|
||||
# Visio autosave temporary files
|
||||
*.~vsd*
|
||||
|
||||
### NotepadPP ###
|
||||
# Notepad++ backups #
|
||||
*.bak
|
||||
|
||||
### Python ###
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
|
||||
# IPython
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/#use-with-ide
|
||||
.pdm.toml
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
### Python Patch ###
|
||||
# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
|
||||
poetry.toml
|
||||
|
||||
# ruff
|
||||
.ruff_cache/
|
||||
|
||||
# LSP config files
|
||||
pyrightconfig.json
|
||||
|
||||
### TortoiseGit ###
|
||||
# Project-level settings
|
||||
/.tgitconfig
|
||||
|
||||
### Vim ###
|
||||
# Swap
|
||||
[._]*.s[a-v][a-z]
|
||||
!*.svg # comment out if you don't need vector files
|
||||
[._]*.sw[a-p]
|
||||
[._]s[a-rt-v][a-z]
|
||||
[._]ss[a-gi-z]
|
||||
[._]sw[a-p]
|
||||
|
||||
# Session
|
||||
Session.vim
|
||||
Sessionx.vim
|
||||
|
||||
# Temporary
|
||||
.netrwhist
|
||||
# Auto-generated tag files
|
||||
tags
|
||||
# Persistent undo
|
||||
[._]*.un~
|
||||
|
||||
### VisualStudioCode ###
|
||||
.vscode/*
|
||||
!.vscode/settings.json
|
||||
!.vscode/tasks.json
|
||||
!.vscode/launch.json
|
||||
!.vscode/extensions.json
|
||||
!.vscode/*.code-snippets
|
||||
|
||||
# Local History for Visual Studio Code
|
||||
.history/
|
||||
|
||||
# Built Visual Studio Code Extensions
|
||||
*.vsix
|
||||
|
||||
### VisualStudioCode Patch ###
|
||||
# Ignore all local history of files
|
||||
.history
|
||||
.ionide
|
||||
|
||||
### Windows ###
|
||||
# Windows thumbnail cache files
|
||||
Thumbs.db
|
||||
Thumbs.db:encryptable
|
||||
ehthumbs.db
|
||||
ehthumbs_vista.db
|
||||
|
||||
# Dump file
|
||||
*.stackdump
|
||||
|
||||
# Folder config file
|
||||
[Dd]esktop.ini
|
||||
|
||||
# Recycle Bin used on file shares
|
||||
$RECYCLE.BIN/
|
||||
|
||||
# Windows Installer files
|
||||
*.cab
|
||||
*.msi
|
||||
*.msix
|
||||
*.msm
|
||||
*.msp
|
||||
|
||||
# Windows shortcuts
|
||||
*.lnk
|
||||
|
||||
-201
@@ -1,201 +0,0 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -1,14 +0,0 @@
|
||||
# from . import timefeatures
|
||||
# from . import precip
|
||||
# from . import metr_la
|
||||
# from .datamodule import DataModule
|
||||
# from .csv_dataset import *
|
||||
# from . import woodside
|
||||
# from . import dow
|
||||
# from . import synthetic
|
||||
# from . import image_completion
|
||||
# from . import copy_task
|
||||
# from . import cont_copy_task
|
||||
# from . import m4
|
||||
# from . import wiki
|
||||
# from . import monash
|
||||
@@ -1,196 +0,0 @@
|
||||
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch import nn
|
||||
import sys
|
||||
|
||||
from src.data.datamodule import DataLoaders
|
||||
from src.data.pred_dataset import *
|
||||
|
||||
DSETS = ['ettm1', 'ettm2', 'etth1', 'etth2', 'electricity',
|
||||
'traffic', 'illness', 'weather', 'exchange'
|
||||
]
|
||||
|
||||
def get_dls(params):
|
||||
|
||||
assert params.dset in DSETS, f"Unrecognized dset (`{params.dset}`). Options include: {DSETS}"
|
||||
if not hasattr(params,'use_time_features'): params.use_time_features = False
|
||||
|
||||
if params.dset == 'ettm1':
|
||||
root_path = '/data/datasets/public/ETDataset/ETT-small/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_ETT_minute,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'ETTm1.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
|
||||
elif params.dset == 'ettm2':
|
||||
root_path = '/data/datasets/public/ETDataset/ETT-small/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_ETT_minute,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'ETTm2.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
elif params.dset == 'etth1':
|
||||
root_path = '/data/datasets/public/ETDataset/ETT-small/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_ETT_hour,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'ETTh1.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
|
||||
elif params.dset == 'etth2':
|
||||
root_path = '/data/datasets/public/ETDataset/ETT-small/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_ETT_hour,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'ETTh2.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
|
||||
elif params.dset == 'electricity':
|
||||
root_path = '/data/datasets/public/electricity/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_Custom,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'electricity.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
elif params.dset == 'traffic':
|
||||
root_path = '/data/datasets/public/traffic/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_Custom,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'traffic.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
elif params.dset == 'weather':
|
||||
root_path = '/data/datasets/public/weather/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_Custom,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'weather.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
elif params.dset == 'illness':
|
||||
root_path = '/data/datasets/public/illness/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_Custom,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'national_illness.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
|
||||
elif params.dset == 'exchange':
|
||||
root_path = '/data/datasets/public/exchange_rate/'
|
||||
size = [params.context_points, 0, params.target_points]
|
||||
dls = DataLoaders(
|
||||
datasetCls=Dataset_Custom,
|
||||
dataset_kwargs={
|
||||
'root_path': root_path,
|
||||
'data_path': 'exchange_rate.csv',
|
||||
'features': params.features,
|
||||
'scale': True,
|
||||
'size': size,
|
||||
'use_time_features': params.use_time_features
|
||||
},
|
||||
batch_size=params.batch_size,
|
||||
workers=params.num_workers,
|
||||
)
|
||||
# dataset is assume to have dimension len x nvars
|
||||
dls.vars, dls.len = dls.train.dataset[0][0].shape[1], params.context_points
|
||||
dls.c = dls.train.dataset[0][1].shape[0]
|
||||
return dls
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
class Params:
|
||||
dset= 'etth2'
|
||||
context_points= 384
|
||||
target_points= 96
|
||||
batch_size= 64
|
||||
num_workers= 8
|
||||
with_ray= False
|
||||
features='M'
|
||||
params = Params
|
||||
dls = get_dls(params)
|
||||
for i, batch in enumerate(dls.valid):
|
||||
print(i, len(batch), batch[0].shape, batch[1].shape)
|
||||
breakpoint()
|
||||
-237
@@ -1,237 +0,0 @@
|
||||
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import os
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from src.models.patchTST import PatchTST
|
||||
from src.learner import Learner, transfer_weights
|
||||
from src.callback.core import *
|
||||
from src.callback.tracking import *
|
||||
from src.callback.patch_mask import *
|
||||
from src.callback.transforms import *
|
||||
from src.metrics import *
|
||||
from src.basics import set_device
|
||||
from datautils import *
|
||||
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
# Pretraining and Finetuning
|
||||
parser.add_argument('--is_finetune', type=int, default=0, help='do finetuning or not')
|
||||
parser.add_argument('--is_linear_probe', type=int, default=0, help='if linear_probe: only finetune the last layer')
|
||||
# Dataset and dataloader
|
||||
parser.add_argument('--dset_finetune', type=str, default='etth1', help='dataset name')
|
||||
parser.add_argument('--context_points', type=int, default=512, help='sequence length')
|
||||
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
|
||||
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
|
||||
parser.add_argument('--num_workers', type=int, default=0, help='number of workers for DataLoader')
|
||||
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
|
||||
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
|
||||
# Patch
|
||||
parser.add_argument('--patch_len', type=int, default=12, help='patch length')
|
||||
parser.add_argument('--stride', type=int, default=12, help='stride between patch')
|
||||
# RevIN
|
||||
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
|
||||
# Model args
|
||||
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
|
||||
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
|
||||
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
|
||||
parser.add_argument('--d_ff', type=int, default=256, help='Tranformer MLP dimension')
|
||||
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
|
||||
parser.add_argument('--head_dropout', type=float, default=0.2, help='head dropout')
|
||||
# Optimization args
|
||||
parser.add_argument('--n_epochs_finetune', type=int, default=20, help='number of finetuning epochs')
|
||||
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
|
||||
# Pretrained model name
|
||||
parser.add_argument('--pretrained_model', type=str, default=None, help='pretrained model name')
|
||||
# model id to keep track of the number of models saved
|
||||
parser.add_argument('--finetuned_model_id', type=int, default=1, help='id of the saved finetuned model')
|
||||
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
print('args:', args)
|
||||
args.save_path = 'saved_models/' + args.dset_finetune + '/masked_patchtst/' + args.model_type + '/'
|
||||
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
|
||||
|
||||
# args.save_finetuned_model = '_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-finetune' + str(args.n_epochs_finetune) + '_mask' + str(args.mask_ratio) + '_model' + str(args.finetuned_model_id)
|
||||
suffix_name = '_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-finetune' + str(args.n_epochs_finetune) + '_model' + str(args.finetuned_model_id)
|
||||
if args.is_finetune: args.save_finetuned_model = args.dset_finetune+'_patchtst_finetuned'+suffix_name
|
||||
elif args.is_linear_probe: args.save_finetuned_model = args.dset_finetune+'_patchtst_linear-probe'+suffix_name
|
||||
else: args.save_finetuned_model = args.dset_finetune+'_patchtst_finetuned'+suffix_name
|
||||
|
||||
# get available GPU devide
|
||||
set_device()
|
||||
|
||||
def get_model(c_in, args, head_type, weight_path=None):
|
||||
"""
|
||||
c_in: number of variables
|
||||
"""
|
||||
# get number of patches
|
||||
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
|
||||
print('number of patches:', num_patch)
|
||||
|
||||
# get model
|
||||
model = PatchTST(c_in=c_in,
|
||||
target_dim=args.target_points,
|
||||
patch_len=args.patch_len,
|
||||
stride=args.stride,
|
||||
num_patch=num_patch,
|
||||
n_layers=args.n_layers,
|
||||
n_heads=args.n_heads,
|
||||
d_model=args.d_model,
|
||||
shared_embedding=True,
|
||||
d_ff=args.d_ff,
|
||||
dropout=args.dropout,
|
||||
head_dropout=args.head_dropout,
|
||||
act='relu',
|
||||
head_type=head_type,
|
||||
res_attention=False
|
||||
)
|
||||
if weight_path: model = transfer_weights(weight_path, model)
|
||||
# print out the model size
|
||||
print('number of model params', sum(p.numel() for p in model.parameters() if p.requires_grad))
|
||||
return model
|
||||
|
||||
|
||||
|
||||
def find_lr(head_type):
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
model = get_model(dls.vars, args, head_type)
|
||||
# transfer weight
|
||||
# weight_path = args.save_path + args.pretrained_model + '.pth'
|
||||
model = transfer_weights(args.pretrained_model, model)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars)] if args.revin else []
|
||||
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
|
||||
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=args.lr,
|
||||
cbs=cbs,
|
||||
)
|
||||
# fit the data to the model
|
||||
suggested_lr = learn.lr_finder()
|
||||
print('suggested_lr', suggested_lr)
|
||||
return suggested_lr
|
||||
|
||||
|
||||
def save_recorders(learn):
|
||||
train_loss = learn.recorder['train_loss']
|
||||
valid_loss = learn.recorder['valid_loss']
|
||||
df = pd.DataFrame(data={'train_loss': train_loss, 'valid_loss': valid_loss})
|
||||
df.to_csv(args.save_path + args.save_finetuned_model + '_losses.csv', float_format='%.6f', index=False)
|
||||
|
||||
|
||||
def finetune_func(lr=args.lr):
|
||||
print('end-to-end finetuning')
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
# get model
|
||||
model = get_model(dls.vars, args, head_type='prediction')
|
||||
# transfer weight
|
||||
# weight_path = args.pretrained_model + '.pth'
|
||||
model = transfer_weights(args.pretrained_model, model)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
|
||||
cbs += [
|
||||
PatchCB(patch_len=args.patch_len, stride=args.stride),
|
||||
SaveModelCB(monitor='valid_loss', fname=args.save_finetuned_model, path=args.save_path)
|
||||
]
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=lr,
|
||||
cbs=cbs,
|
||||
metrics=[mse]
|
||||
)
|
||||
# fit the data to the model
|
||||
#learn.fit_one_cycle(n_epochs=args.n_epochs_finetune, lr_max=lr)
|
||||
learn.fine_tune(n_epochs=args.n_epochs_finetune, base_lr=lr, freeze_epochs=10)
|
||||
save_recorders(learn)
|
||||
|
||||
|
||||
def linear_probe_func(lr=args.lr):
|
||||
print('linear probing')
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
# get model
|
||||
model = get_model(dls.vars, args, head_type='prediction')
|
||||
# transfer weight
|
||||
# weight_path = args.save_path + args.pretrained_model + '.pth'
|
||||
model = transfer_weights(args.pretrained_model, model)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
|
||||
cbs += [
|
||||
PatchCB(patch_len=args.patch_len, stride=args.stride),
|
||||
SaveModelCB(monitor='valid_loss', fname=args.save_finetuned_model, path=args.save_path)
|
||||
]
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=lr,
|
||||
cbs=cbs,
|
||||
metrics=[mse]
|
||||
)
|
||||
# fit the data to the model
|
||||
learn.linear_probe(n_epochs=args.n_epochs_finetune, base_lr=lr)
|
||||
save_recorders(learn)
|
||||
|
||||
|
||||
def test_func(weight_path):
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
model = get_model(dls.vars, args, head_type='prediction').to('cuda')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars, denorm=True)] if args.revin else []
|
||||
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
|
||||
learn = Learner(dls, model,cbs=cbs)
|
||||
out = learn.test(dls.test, weight_path=weight_path+'.pth', scores=[mse,mae]) # out: a list of [pred, targ, score]
|
||||
print('score:', out[2])
|
||||
# save results
|
||||
pd.DataFrame(np.array(out[2]).reshape(1,-1), columns=['mse','mae']).to_csv(args.save_path + args.save_finetuned_model + '_acc.csv', float_format='%.6f', index=False)
|
||||
return out
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
if args.is_finetune:
|
||||
args.dset = args.dset_finetune
|
||||
# Finetune
|
||||
suggested_lr = find_lr(head_type='prediction')
|
||||
finetune_func(suggested_lr)
|
||||
print('finetune completed')
|
||||
# Test
|
||||
out = test_func(args.save_path+args.save_finetuned_model)
|
||||
print('----------- Complete! -----------')
|
||||
|
||||
elif args.is_linear_probe:
|
||||
args.dset = args.dset_finetune
|
||||
# Finetune
|
||||
suggested_lr = find_lr(head_type='prediction')
|
||||
linear_probe_func(suggested_lr)
|
||||
print('finetune completed')
|
||||
# Test
|
||||
out = test_func(args.save_path+args.save_finetuned_model)
|
||||
print('----------- Complete! -----------')
|
||||
|
||||
else:
|
||||
args.dset = args.dset_finetune
|
||||
weight_path = args.save_path+args.dset_finetune+'_patchtst_finetuned'+suffix_name
|
||||
# Test
|
||||
out = test_func(weight_path)
|
||||
print('----------- Complete! -----------')
|
||||
|
||||
|
||||
-153
@@ -1,153 +0,0 @@
|
||||
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import os
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from src.models.patchTST import PatchTST
|
||||
from src.learner import Learner, transfer_weights
|
||||
from src.callback.tracking import *
|
||||
from src.callback.patch_mask import *
|
||||
from src.callback.transforms import *
|
||||
from src.metrics import *
|
||||
from src.basics import set_device
|
||||
from datautils import *
|
||||
|
||||
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
# Dataset and dataloader
|
||||
parser.add_argument('--dset_pretrain', type=str, default='etth1', help='dataset name')
|
||||
parser.add_argument('--context_points', type=int, default=512, help='sequence length')
|
||||
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
|
||||
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
|
||||
parser.add_argument('--num_workers', type=int, default=0, help='number of workers for DataLoader')
|
||||
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
|
||||
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
|
||||
# Patch
|
||||
parser.add_argument('--patch_len', type=int, default=12, help='patch length')
|
||||
parser.add_argument('--stride', type=int, default=12, help='stride between patch')
|
||||
# RevIN
|
||||
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
|
||||
# Model args
|
||||
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
|
||||
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
|
||||
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
|
||||
parser.add_argument('--d_ff', type=int, default=512, help='Tranformer MLP dimension')
|
||||
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
|
||||
parser.add_argument('--head_dropout', type=float, default=0.2, help='head dropout')
|
||||
# Pretrain mask
|
||||
parser.add_argument('--mask_ratio', type=float, default=0.4, help='masking ratio for the input')
|
||||
# Optimization args
|
||||
parser.add_argument('--n_epochs_pretrain', type=int, default=10, help='number of pre-training epochs')
|
||||
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
|
||||
# model id to keep track of the number of models saved
|
||||
parser.add_argument('--pretrained_model_id', type=int, default=1, help='id of the saved pretrained model')
|
||||
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
print('args:', args)
|
||||
args.save_pretrained_model = 'patchtst_pretrained_cw'+str(args.context_points)+'_patch'+str(args.patch_len) + '_stride'+str(args.stride) + '_epochs-pretrain' + str(args.n_epochs_pretrain) + '_mask' + str(args.mask_ratio) + '_model' + str(args.pretrained_model_id)
|
||||
args.save_path = 'saved_models/' + args.dset_pretrain + '/masked_patchtst/' + args.model_type + '/'
|
||||
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
|
||||
|
||||
|
||||
# get available GPU devide
|
||||
set_device()
|
||||
|
||||
|
||||
def get_model(c_in, args):
|
||||
"""
|
||||
c_in: number of variables
|
||||
"""
|
||||
# get number of patches
|
||||
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
|
||||
print('number of patches:', num_patch)
|
||||
|
||||
# get model
|
||||
model = PatchTST(c_in=c_in,
|
||||
target_dim=args.target_points,
|
||||
patch_len=args.patch_len,
|
||||
stride=args.stride,
|
||||
num_patch=num_patch,
|
||||
n_layers=args.n_layers,
|
||||
n_heads=args.n_heads,
|
||||
d_model=args.d_model,
|
||||
shared_embedding=True,
|
||||
d_ff=args.d_ff,
|
||||
dropout=args.dropout,
|
||||
head_dropout=args.head_dropout,
|
||||
act='relu',
|
||||
head_type='pretrain',
|
||||
res_attention=False
|
||||
)
|
||||
# print out the model size
|
||||
print('number of model params', sum(p.numel() for p in model.parameters() if p.requires_grad))
|
||||
return model
|
||||
|
||||
|
||||
def find_lr():
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
model = get_model(dls.vars, args)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars, denorm=False)] if args.revin else []
|
||||
cbs += [PatchMaskCB(patch_len=args.patch_len, stride=args.stride, mask_ratio=args.mask_ratio)]
|
||||
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=args.lr,
|
||||
cbs=cbs,
|
||||
)
|
||||
# fit the data to the model
|
||||
suggested_lr = learn.lr_finder()
|
||||
print('suggested_lr', suggested_lr)
|
||||
return suggested_lr
|
||||
|
||||
|
||||
def pretrain_func(lr=args.lr):
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
# get model
|
||||
model = get_model(dls.vars, args)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars, denorm=False)] if args.revin else []
|
||||
cbs += [
|
||||
PatchMaskCB(patch_len=args.patch_len, stride=args.stride, mask_ratio=args.mask_ratio),
|
||||
SaveModelCB(monitor='valid_loss', fname=args.save_pretrained_model,
|
||||
path=args.save_path)
|
||||
]
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=lr,
|
||||
cbs=cbs,
|
||||
#metrics=[mse]
|
||||
)
|
||||
# fit the data to the model
|
||||
learn.fit_one_cycle(n_epochs=args.n_epochs_pretrain, lr_max=lr)
|
||||
|
||||
train_loss = learn.recorder['train_loss']
|
||||
valid_loss = learn.recorder['valid_loss']
|
||||
df = pd.DataFrame(data={'train_loss': train_loss, 'valid_loss': valid_loss})
|
||||
df.to_csv(args.save_path + args.save_pretrained_model + '_losses.csv', float_format='%.6f', index=False)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
args.dset = args.dset_pretrain
|
||||
suggested_lr = find_lr()
|
||||
# Pretrain
|
||||
pretrain_func(suggested_lr)
|
||||
print('pretraining completed')
|
||||
|
||||
|
||||
-162
@@ -1,162 +0,0 @@
|
||||
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import os
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from src.models.patchTST import PatchTST
|
||||
from src.learner import Learner
|
||||
from src.callback.core import *
|
||||
from src.callback.tracking import *
|
||||
from src.callback.scheduler import *
|
||||
from src.callback.patch_mask import *
|
||||
from src.callback.transforms import *
|
||||
from src.metrics import *
|
||||
from datautils import get_dls
|
||||
|
||||
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
# Dataset and dataloader
|
||||
parser.add_argument('--dset', type=str, default='etth1', help='dataset name')
|
||||
parser.add_argument('--context_points', type=int, default=336, help='sequence length')
|
||||
parser.add_argument('--target_points', type=int, default=96, help='forecast horizon')
|
||||
parser.add_argument('--batch_size', type=int, default=64, help='batch size')
|
||||
parser.add_argument('--num_workers', type=int, default=1, help='number of workers for DataLoader')
|
||||
parser.add_argument('--scaler', type=str, default='standard', help='scale the input data')
|
||||
parser.add_argument('--features', type=str, default='M', help='for multivariate model or univariate model')
|
||||
parser.add_argument('--use_time_features', type=int, default=0, help='whether to use time features or not')
|
||||
# Patch
|
||||
parser.add_argument('--patch_len', type=int, default=32, help='patch length')
|
||||
parser.add_argument('--stride', type=int, default=16, help='stride between patch')
|
||||
# RevIN
|
||||
parser.add_argument('--revin', type=int, default=1, help='reversible instance normalization')
|
||||
# Model args
|
||||
parser.add_argument('--n_layers', type=int, default=3, help='number of Transformer layers')
|
||||
parser.add_argument('--n_heads', type=int, default=16, help='number of Transformer heads')
|
||||
parser.add_argument('--d_model', type=int, default=128, help='Transformer d_model')
|
||||
parser.add_argument('--d_ff', type=int, default=256, help='Tranformer MLP dimension')
|
||||
parser.add_argument('--dropout', type=float, default=0.2, help='Transformer dropout')
|
||||
parser.add_argument('--head_dropout', type=float, default=0, help='head dropout')
|
||||
# Optimization args
|
||||
parser.add_argument('--n_epochs', type=int, default=20, help='number of training epochs')
|
||||
parser.add_argument('--lr', type=float, default=1e-4, help='learning rate')
|
||||
# model id to keep track of the number of models saved
|
||||
parser.add_argument('--model_id', type=int, default=1, help='id of the saved model')
|
||||
parser.add_argument('--model_type', type=str, default='based_model', help='for multivariate model or univariate model')
|
||||
# training
|
||||
parser.add_argument('--is_train', type=int, default=1, help='training the model')
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
print('args:', args)
|
||||
args.save_model_name = 'patchtst_supervised'+'_cw'+str(args.context_points)+'_tw'+str(args.target_points) + '_patch'+str(args.patch_len) + '_stride'+str(args.stride)+'_epochs'+str(args.n_epochs) + '_model' + str(args.model_id)
|
||||
args.save_path = 'saved_models/' + args.dset + '/patchtst_supervised/' + args.model_type + '/'
|
||||
if not os.path.exists(args.save_path): os.makedirs(args.save_path)
|
||||
|
||||
|
||||
def get_model(c_in, args):
|
||||
"""
|
||||
c_in: number of input variables
|
||||
"""
|
||||
# get number of patches
|
||||
num_patch = (max(args.context_points, args.patch_len)-args.patch_len) // args.stride + 1
|
||||
print('number of patches:', num_patch)
|
||||
|
||||
# get model
|
||||
model = PatchTST(c_in=c_in,
|
||||
target_dim=args.target_points,
|
||||
patch_len=args.patch_len,
|
||||
stride=args.stride,
|
||||
num_patch=num_patch,
|
||||
n_layers=args.n_layers,
|
||||
n_heads=args.n_heads,
|
||||
d_model=args.d_model,
|
||||
shared_embedding=True,
|
||||
d_ff=args.d_ff,
|
||||
dropout=args.dropout,
|
||||
head_dropout=args.head_dropout,
|
||||
act='relu',
|
||||
head_type='prediction',
|
||||
res_attention=False
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def find_lr():
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
model = get_model(dls.vars, args)
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars)] if args.revin else []
|
||||
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
|
||||
# define learner
|
||||
learn = Learner(dls, model, loss_func, cbs=cbs)
|
||||
# fit the data to the model
|
||||
return learn.lr_finder()
|
||||
|
||||
|
||||
def train_func(lr=args.lr):
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
print('in out', dls.vars, dls.c, dls.len)
|
||||
|
||||
# get model
|
||||
model = get_model(dls.vars, args)
|
||||
|
||||
# get loss
|
||||
loss_func = torch.nn.MSELoss(reduction='mean')
|
||||
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars)] if args.revin else []
|
||||
cbs += [
|
||||
PatchCB(patch_len=args.patch_len, stride=args.stride),
|
||||
SaveModelCB(monitor='valid_loss', fname=args.save_model_name,
|
||||
path=args.save_path )
|
||||
]
|
||||
|
||||
# define learner
|
||||
learn = Learner(dls, model,
|
||||
loss_func,
|
||||
lr=lr,
|
||||
cbs=cbs,
|
||||
metrics=[mse]
|
||||
)
|
||||
|
||||
# fit the data to the model
|
||||
learn.fit_one_cycle(n_epochs=args.n_epochs, lr_max=lr, pct_start=0.2)
|
||||
|
||||
|
||||
def test_func():
|
||||
weight_path = args.save_path + args.save_model_name + '.pth'
|
||||
# get dataloader
|
||||
dls = get_dls(args)
|
||||
model = get_model(dls.vars, args)
|
||||
#model = torch.load(weight_path)
|
||||
# get callbacks
|
||||
cbs = [RevInCB(dls.vars)] if args.revin else []
|
||||
cbs += [PatchCB(patch_len=args.patch_len, stride=args.stride)]
|
||||
learn = Learner(dls, model,cbs=cbs)
|
||||
out = learn.test(dls.test, weight_path=weight_path, scores=[mse,mae]) # out: a list of [pred, targ, score_values]
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
if args.is_train: # training mode
|
||||
suggested_lr = find_lr()
|
||||
print('suggested lr:', suggested_lr)
|
||||
train_func(suggested_lr)
|
||||
else: # testing mode
|
||||
out = test_func()
|
||||
print('score:', out[2])
|
||||
print('shape:', out[0].shape)
|
||||
|
||||
print('----------- Complete! -----------')
|
||||
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
|
||||
import torch
|
||||
|
||||
import collections
|
||||
from collections import OrderedDict
|
||||
|
||||
class GetAttr:
|
||||
|
||||
"Inherit from this to have all attr accesses in `self._xtra` passed down to `self.default`"
|
||||
_default='default'
|
||||
def _component_attr_filter(self,k):
|
||||
if k.startswith('__') or k in ('_xtra',self._default): return False
|
||||
xtra = getattr(self,'_xtra',None)
|
||||
return xtra is None or k in xtra
|
||||
|
||||
def _dir(self):
|
||||
return [k for k in dir(getattr(self,self._default)) if self._component_attr_filter(k)]
|
||||
|
||||
def __getattr__(self, k):
|
||||
if self._component_attr_filter(k):
|
||||
attr = getattr(self, self._default, None)
|
||||
if attr is not None: return getattr(attr,k)
|
||||
# raise AttributeError(k)
|
||||
|
||||
def __dir__(self):
|
||||
return custom_dir(self,self._dir())
|
||||
|
||||
# def __getstate__(self): return self.__dict__
|
||||
def __setstate__(self,data):
|
||||
self.__dict__.update(data)
|
||||
|
||||
|
||||
|
||||
def get_device(use_cuda=True, device_id=None, usage=5):
|
||||
"Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU"
|
||||
if not torch.cuda.is_available():
|
||||
use_cuda = False
|
||||
else:
|
||||
if device_id is None:
|
||||
device_ids = get_available_cuda(usage=usage)
|
||||
device_id = device_ids[0] # get the first available device
|
||||
torch.cuda.set_device(device_id)
|
||||
return torch.device(torch.cuda.current_device()) if use_cuda else torch.device('cpu')
|
||||
|
||||
|
||||
def set_device(usage=5):
|
||||
"set the device that has usage < default usage "
|
||||
device_ids = get_available_cuda(usage=usage)
|
||||
torch.cuda.set_device(device_ids[0]) # get the first available device
|
||||
|
||||
|
||||
def default_device(use_cuda=True):
|
||||
"Return or set default device; `use_cuda`: None - CUDA if available; True - error if not available; False - CPU"
|
||||
if not torch.cuda.is_available():
|
||||
use_cuda = False
|
||||
return torch.device(torch.cuda.current_device()) if use_cuda else torch.device('cpu')
|
||||
|
||||
|
||||
def get_available_cuda(usage=10):
|
||||
if not torch.cuda.is_available(): return
|
||||
# collect available cuda devices, only collect devices that has less that 'usage' percent
|
||||
device_ids = []
|
||||
for device in range(torch.cuda.device_count()):
|
||||
if torch.cuda.utilization(device) < usage: device_ids.append(device)
|
||||
return device_ids
|
||||
|
||||
|
||||
|
||||
def to_device(b, device=None, non_blocking=False):
|
||||
"""
|
||||
Recursively put `b` on `device`
|
||||
components of b are torch tensors
|
||||
"""
|
||||
if device is None:
|
||||
device = default_device(use_cuda=True)
|
||||
|
||||
if isinstance(b, dict):
|
||||
return {key: to_device(val, device) for key, val in b.items()}
|
||||
|
||||
if isinstance(b, (list, tuple)):
|
||||
return type(b)(to_device(o, device) for o in b)
|
||||
|
||||
return b.to(device, non_blocking=non_blocking)
|
||||
|
||||
|
||||
def to_numpy(b):
|
||||
"""
|
||||
Components of b are torch tensors
|
||||
"""
|
||||
if isinstance(b, dict):
|
||||
return {key: to_numpy(val) for key, val in b.items()}
|
||||
|
||||
if isinstance(b, (list, tuple)):
|
||||
return type(b)(to_numpy(o) for o in b)
|
||||
|
||||
return b.detach().cpu().numpy()
|
||||
|
||||
-11
@@ -1,11 +0,0 @@
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
-90
@@ -1,90 +0,0 @@
|
||||
|
||||
__all__ = ['Callback', 'SetupLearnerCB', 'GetPredictionsCB', 'GetTestCB' ]
|
||||
|
||||
|
||||
"""
|
||||
Callback lists:
|
||||
> before_fit
|
||||
- before_epoch
|
||||
+ before_epoch_train
|
||||
~ before_batch_train
|
||||
~ after_batch_train
|
||||
+ after_epoch_train
|
||||
|
||||
+ before_epoch_valid
|
||||
~ before_batch_valid
|
||||
~ after_batch_valid
|
||||
+ after_epoch_valid
|
||||
- after_epoch
|
||||
> after_fit
|
||||
|
||||
- before_predict
|
||||
~ before_batch_predict
|
||||
~ after_batch_predict
|
||||
- after_predict
|
||||
|
||||
"""
|
||||
|
||||
from ..basics import *
|
||||
import torch
|
||||
|
||||
DTYPE = torch.float32
|
||||
|
||||
class Callback(GetAttr):
|
||||
_default='learner'
|
||||
|
||||
|
||||
class SetupLearnerCB(Callback):
|
||||
def __init__(self):
|
||||
self.device = default_device(use_cuda=True)
|
||||
|
||||
def before_batch_train(self): self._to_device()
|
||||
def before_batch_valid(self): self._to_device()
|
||||
def before_batch_predict(self): self._to_device()
|
||||
def before_batch_test(self): self._to_device()
|
||||
|
||||
def _to_device(self):
|
||||
batch = to_device(self.batch, self.device)
|
||||
if self.n_inp > 1: xb, yb = batch
|
||||
else: xb, yb = batch, None
|
||||
self.learner.batch = xb, yb
|
||||
|
||||
def before_fit(self):
|
||||
"Set model to cuda before training"
|
||||
self.learner.model.to(self.device)
|
||||
self.learner.device = self.device
|
||||
|
||||
|
||||
class GetPredictionsCB(Callback):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def before_predict(self):
|
||||
self.preds = []
|
||||
|
||||
def after_batch_predict(self):
|
||||
# append the prediction after each forward batch
|
||||
self.preds.append(self.pred)
|
||||
|
||||
def after_predict(self):
|
||||
self.preds = torch.concat(self.preds)#.detach().cpu().numpy()
|
||||
|
||||
|
||||
|
||||
class GetTestCB(Callback):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def before_test(self):
|
||||
self.preds, self.targets = [], []
|
||||
|
||||
def after_batch_test(self):
|
||||
# append the prediction after each forward batch
|
||||
self.preds.append(self.pred)
|
||||
self.targets.append(self.yb)
|
||||
|
||||
def after_test(self):
|
||||
self.preds = torch.concat(self.preds)#.detach().cpu().numpy()
|
||||
self.targets = torch.concat(self.targets)#.detach().cpu().numpy()
|
||||
|
||||
|
||||
-173
@@ -1,173 +0,0 @@
|
||||
from .core import Callback
|
||||
import torch
|
||||
from torch.utils.data import DistributedSampler, DataLoader, SequentialSampler
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from typing import Optional, Dict, Any
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DistributedTrainer(Callback):
|
||||
"Wrap `model` in `DistributedDataParallel` and `dls` in `DistributedDL`"
|
||||
def __init__(self,
|
||||
local_rank,
|
||||
world_size,
|
||||
sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`
|
||||
**kwargs
|
||||
):
|
||||
self.local_rank = local_rank
|
||||
self.world_size = world_size
|
||||
self.sync_bn = sync_bn
|
||||
self.kwargs = kwargs
|
||||
super().__init__()
|
||||
|
||||
def before_fit(self):
|
||||
self.learner.model = self.prepare_model(
|
||||
torch.nn.SyncBatchNorm.convert_sync_batchnorm(self.model) if self.sync_bn else self.model,
|
||||
ddp_kwargs=self.kwargs
|
||||
)
|
||||
self.old_train_dl = self.dls.train
|
||||
self.old_valid_dl = self.dls.valid
|
||||
|
||||
self.learner.dls.train = self._wrap_dl(self.dls.train)
|
||||
self.learner.dls.valid = self._wrap_dl(self.dls.valid)
|
||||
|
||||
def _wrap_dl(self, dl):
|
||||
return dl if isinstance(dl, DistributedDL) else self.prepare_data_loader(dl)
|
||||
|
||||
|
||||
def after_fit(self):
|
||||
self.learner.model = self.learner.model.module
|
||||
self.learner.dls.train = self.old_train_dl
|
||||
self.learner.dls.valid = self.old_valid_dl
|
||||
|
||||
def prepare_model(self,
|
||||
model: torch.nn.Module,
|
||||
move_to_device: bool = True,
|
||||
wrap_ddp: bool = True,
|
||||
ddp_kwargs: Optional[Dict[str, Any]] = None) -> torch.nn.Module:
|
||||
"""Prepares the model for distributed execution.
|
||||
Args:
|
||||
model (torch.nn.Module): A torch model to prepare.
|
||||
move_to_device (bool): Whether to move the model to the correct
|
||||
device. If set to False, the model needs to manually be moved
|
||||
to the correct device.
|
||||
wrap_ddp (bool): Whether to wrap models in
|
||||
``DistributedDataParallel``.
|
||||
ddp_kwargs (Dict[str, Any]): Args to pass into
|
||||
``DistributedDataParallel`` initialization if ``wrap_ddp`` is
|
||||
set to True.
|
||||
"""
|
||||
ddp_kwargs = ddp_kwargs or {}
|
||||
|
||||
rank = self.local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
# device = get_device()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
if move_to_device:
|
||||
logger.info(f"Moving model to device: {device}")
|
||||
model = model.to(device)
|
||||
if wrap_ddp and self.world_size > 1:
|
||||
logger.info("Wrapping provided model in DDP.")
|
||||
if torch.cuda.is_available():
|
||||
model = DistributedDataParallel(
|
||||
model, device_ids=[rank], output_device=rank, **ddp_kwargs)
|
||||
else:
|
||||
model = DistributedDataParallel(model, **ddp_kwargs)
|
||||
|
||||
return model
|
||||
|
||||
def prepare_data_loader(self,
|
||||
data_loader: torch.utils.data.DataLoader,
|
||||
add_dist_sampler: bool = True,
|
||||
move_to_device: bool = True) -> torch.utils.data.DataLoader:
|
||||
"""
|
||||
Prepares DataLoader for distributed execution.
|
||||
|
||||
This allows you to use the same exact code regardless of number of
|
||||
workers or the device type being used (CPU, GPU).
|
||||
|
||||
Args:
|
||||
data_loader (torch.utils.data.DataLoader): The DataLoader to
|
||||
prepare.
|
||||
add_dist_sampler (bool): Whether to add a DistributedSampler to
|
||||
the provided DataLoader.
|
||||
move_to_device (bool): If set, automatically move the data
|
||||
returned by the data loader to the correct device.
|
||||
"""
|
||||
|
||||
# Only add Distributed Sampler if the following conditions hold:
|
||||
# 1. More than one training worker is being used.
|
||||
# 2. A DistributedSampler has not already been added by the user.
|
||||
# 3. The dataset is not an IterableDataset. Samplers do not worker with
|
||||
# IterableDatasets.
|
||||
def with_sampler(loader):
|
||||
# Automatically set the DistributedSampler
|
||||
|
||||
# If using a sampler, the shuffle attribute in the
|
||||
# DataLoader must be set to False.
|
||||
# Instead the shuffling is determined by the shuffle attribute
|
||||
# in the DistributedSampler.
|
||||
# We identify if shuffling is enabled in the passed in
|
||||
# DataLoader by seeing if the sampler for the DataLoader is a
|
||||
# SequentialSampler.
|
||||
shuffle = not isinstance(loader.sampler, SequentialSampler)
|
||||
|
||||
data_loader_args = {
|
||||
"dataset": loader.dataset,
|
||||
"batch_size": loader.batch_size,
|
||||
"shuffle": False,
|
||||
"num_workers": loader.num_workers,
|
||||
"collate_fn": loader.collate_fn,
|
||||
"pin_memory": loader.pin_memory,
|
||||
"drop_last": loader.drop_last,
|
||||
"timeout": loader.timeout,
|
||||
"worker_init_fn": loader.worker_init_fn,
|
||||
"sampler": DistributedSampler(loader.dataset, shuffle=shuffle)
|
||||
}
|
||||
return DataLoader(**data_loader_args)
|
||||
|
||||
data_loader = with_sampler(data_loader)
|
||||
|
||||
if move_to_device:
|
||||
if torch.cuda.is_available():
|
||||
rank = self.local_rank
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
data_loader = DistributedDL(data_loader, device)
|
||||
|
||||
return data_loader
|
||||
|
||||
|
||||
class DistributedDL(DataLoader):
|
||||
def __init__(self, base_dataloader: DataLoader, device: torch.device):
|
||||
|
||||
self.__dict__.update(getattr(base_dataloader, "__dict__", {}))
|
||||
self.dataloader = base_dataloader
|
||||
self.device = device
|
||||
|
||||
def _move_to_device(self, item):
|
||||
def try_move_device(i):
|
||||
try:
|
||||
i = i.to(self.device)
|
||||
except AttributeError:
|
||||
logger.debug(f"Item {i} cannot be moved to device "
|
||||
f"{self.device}.")
|
||||
return i
|
||||
|
||||
return tuple(try_move_device(i) for i in item)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataloader)
|
||||
|
||||
def __iter__(self):
|
||||
iterator = iter(self.dataloader)
|
||||
|
||||
for item in iterator:
|
||||
yield self._move_to_device(item)
|
||||
-176
@@ -1,176 +0,0 @@
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .core import Callback
|
||||
|
||||
# Cell
|
||||
class PatchCB(Callback):
|
||||
|
||||
def __init__(self, patch_len, stride ):
|
||||
"""
|
||||
Callback used to perform patching on the batch input data
|
||||
Args:
|
||||
patch_len: patch length
|
||||
stride: stride
|
||||
"""
|
||||
self.patch_len = patch_len
|
||||
self.stride = stride
|
||||
|
||||
def before_forward(self): self.set_patch()
|
||||
|
||||
def set_patch(self):
|
||||
"""
|
||||
take xb from learner and convert to patch: [bs x seq_len x n_vars] -> [bs x num_patch x n_vars x patch_len]
|
||||
"""
|
||||
xb_patch, num_patch = create_patch(self.xb, self.patch_len, self.stride) # xb: [bs x seq_len x n_vars]
|
||||
# learner get the transformed input
|
||||
self.learner.xb = xb_patch # xb_patch: [bs x num_patch x n_vars x patch_len]
|
||||
|
||||
|
||||
class PatchMaskCB(Callback):
|
||||
def __init__(self, patch_len, stride, mask_ratio,
|
||||
mask_when_pred:bool=False):
|
||||
"""
|
||||
Callback used to perform the pretext task of reconstruct the original data after a binary mask has been applied.
|
||||
Args:
|
||||
patch_len: patch length
|
||||
stride: stride
|
||||
mask_ratio: mask ratio
|
||||
"""
|
||||
self.patch_len = patch_len
|
||||
self.stride = stride
|
||||
self.mask_ratio = mask_ratio
|
||||
|
||||
def before_fit(self):
|
||||
# overwrite the predefined loss function
|
||||
self.learner.loss_func = self._loss
|
||||
device = self.learner.device
|
||||
|
||||
def before_forward(self): self.patch_masking()
|
||||
|
||||
def patch_masking(self):
|
||||
"""
|
||||
xb: [bs x seq_len x n_vars] -> [bs x num_patch x n_vars x patch_len]
|
||||
"""
|
||||
xb_patch, num_patch = create_patch(self.xb, self.patch_len, self.stride) # xb_patch: [bs x num_patch x n_vars x patch_len]
|
||||
xb_mask, _, self.mask, _ = random_masking(xb_patch, self.mask_ratio) # xb_mask: [bs x num_patch x n_vars x patch_len]
|
||||
self.mask = self.mask.bool() # mask: [bs x num_patch x n_vars]
|
||||
self.learner.xb = xb_mask # learner.xb: masked 4D tensor
|
||||
self.learner.yb = xb_patch # learner.yb: non-masked 4d tensor
|
||||
|
||||
def _loss(self, preds, target):
|
||||
"""
|
||||
preds: [bs x num_patch x n_vars x patch_len]
|
||||
targets: [bs x num_patch x n_vars x patch_len]
|
||||
"""
|
||||
loss = (preds - target) ** 2
|
||||
loss = loss.mean(dim=-1)
|
||||
loss = (loss * self.mask).sum() / self.mask.sum()
|
||||
return loss
|
||||
|
||||
|
||||
def create_patch(xb, patch_len, stride):
|
||||
"""
|
||||
xb: [bs x seq_len x n_vars]
|
||||
"""
|
||||
seq_len = xb.shape[1]
|
||||
num_patch = (max(seq_len, patch_len)-patch_len) // stride + 1
|
||||
tgt_len = patch_len + stride*(num_patch-1)
|
||||
s_begin = seq_len - tgt_len
|
||||
|
||||
xb = xb[:, s_begin:, :] # xb: [bs x tgt_len x nvars]
|
||||
xb = xb.unfold(dimension=1, size=patch_len, step=stride) # xb: [bs x num_patch x n_vars x patch_len]
|
||||
return xb, num_patch
|
||||
|
||||
|
||||
class Patch(nn.Module):
|
||||
def __init__(self,seq_len, patch_len, stride):
|
||||
super().__init__()
|
||||
self.seq_len = seq_len
|
||||
self.patch_len = patch_len
|
||||
self.stride = stride
|
||||
self.num_patch = (max(seq_len, patch_len)-patch_len) // stride + 1
|
||||
tgt_len = patch_len + stride*(self.num_patch-1)
|
||||
self.s_begin = seq_len - tgt_len
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x seq_len x n_vars]
|
||||
"""
|
||||
x = x[:, self.s_begin:, :]
|
||||
x = x.unfold(dimension=1, size=self.patch_len, step=self.stride) # xb: [bs x num_patch x n_vars x patch_len]
|
||||
return x
|
||||
|
||||
|
||||
def random_masking(xb, mask_ratio):
|
||||
# xb: [bs x num_patch x n_vars x patch_len]
|
||||
bs, L, nvars, D = xb.shape
|
||||
x = xb.clone()
|
||||
|
||||
len_keep = int(L * (1 - mask_ratio))
|
||||
|
||||
noise = torch.rand(bs, L, nvars,device=xb.device) # noise in [0, 1], bs x L x nvars
|
||||
|
||||
# sort noise for each sample
|
||||
ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove
|
||||
ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore: [bs x L x nvars]
|
||||
|
||||
# keep the first subset
|
||||
ids_keep = ids_shuffle[:, :len_keep, :] # ids_keep: [bs x len_keep x nvars]
|
||||
x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, 1, D)) # x_kept: [bs x len_keep x nvars x patch_len]
|
||||
|
||||
# removed x
|
||||
x_removed = torch.zeros(bs, L-len_keep, nvars, D, device=xb.device) # x_removed: [bs x (L-len_keep) x nvars x patch_len]
|
||||
x_ = torch.cat([x_kept, x_removed], dim=1) # x_: [bs x L x nvars x patch_len]
|
||||
|
||||
# combine the kept part and the removed one
|
||||
x_masked = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1,1,1,D)) # x_masked: [bs x num_patch x nvars x patch_len]
|
||||
|
||||
# generate the binary mask: 0 is keep, 1 is remove
|
||||
mask = torch.ones([bs, L, nvars], device=x.device) # mask: [bs x num_patch x nvars]
|
||||
mask[:, :len_keep, :] = 0
|
||||
# unshuffle to get the binary mask
|
||||
mask = torch.gather(mask, dim=1, index=ids_restore) # [bs x num_patch x nvars]
|
||||
return x_masked, x_kept, mask, ids_restore
|
||||
|
||||
|
||||
def random_masking_3D(xb, mask_ratio):
|
||||
# xb: [bs x num_patch x dim]
|
||||
bs, L, D = xb.shape
|
||||
x = xb.clone()
|
||||
|
||||
len_keep = int(L * (1 - mask_ratio))
|
||||
|
||||
noise = torch.rand(bs, L, device=xb.device) # noise in [0, 1], bs x L
|
||||
|
||||
# sort noise for each sample
|
||||
ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove
|
||||
ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore: [bs x L]
|
||||
|
||||
# keep the first subset
|
||||
ids_keep = ids_shuffle[:, :len_keep] # ids_keep: [bs x len_keep]
|
||||
x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) # x_kept: [bs x len_keep x dim]
|
||||
|
||||
# removed x
|
||||
x_removed = torch.zeros(bs, L-len_keep, D, device=xb.device) # x_removed: [bs x (L-len_keep) x dim]
|
||||
x_ = torch.cat([x_kept, x_removed], dim=1) # x_: [bs x L x dim]
|
||||
|
||||
# combine the kept part and the removed one
|
||||
x_masked = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1,1,D)) # x_masked: [bs x num_patch x dim]
|
||||
|
||||
# generate the binary mask: 0 is keep, 1 is remove
|
||||
mask = torch.ones([bs, L], device=x.device) # mask: [bs x num_patch]
|
||||
mask[:, :len_keep] = 0
|
||||
# unshuffle to get the binary mask
|
||||
mask = torch.gather(mask, dim=1, index=ids_restore) # [bs x num_patch]
|
||||
return x_masked, x_kept, mask, ids_restore
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
bs, L, nvars, D = 2,20,4,5
|
||||
xb = torch.randn(bs, L, nvars, D)
|
||||
xb_mask, mask, ids_restore = create_mask(xb, mask_ratio=0.5)
|
||||
breakpoint()
|
||||
|
||||
|
||||
-212
@@ -1,212 +0,0 @@
|
||||
|
||||
__all__ = ['OneCycleLR', 'LRFinderCB', 'LinearLR', 'ExponentialLR']
|
||||
|
||||
from cmath import inf
|
||||
from ..basics import *
|
||||
from .core import Callback
|
||||
from torch.optim import lr_scheduler
|
||||
from torch.optim.lr_scheduler import _LRScheduler
|
||||
|
||||
|
||||
|
||||
class OneCycleLR(Callback):
|
||||
def __init__(self, lr_max=None,
|
||||
total_steps=None,
|
||||
steps_per_epoch=None,
|
||||
pct_start=0.3,
|
||||
anneal_strategy='cos',
|
||||
cycle_momentum=True,
|
||||
base_momentum=0.85,
|
||||
max_momentum=0.95,
|
||||
div_factor=25.,
|
||||
final_div_factor=1e4,
|
||||
three_phase=False,
|
||||
last_epoch=-1,
|
||||
verbose=False):
|
||||
|
||||
super().__init__()
|
||||
self.lr_max = lr_max if lr_max else self.lr
|
||||
self.total_steps, self.steps_per_epoch = total_steps, steps_per_epoch
|
||||
self.pct_start = pct_start
|
||||
self.anneal_strategy, self.cycle_momentum = anneal_strategy, cycle_momentum
|
||||
self.base_momentum, self.max_momentum = base_momentum, max_momentum
|
||||
self.div_factor, self.final_div_factor = div_factor, final_div_factor
|
||||
self.three_phase = three_phase
|
||||
self.last_epoch = last_epoch
|
||||
self.verbose = verbose
|
||||
|
||||
|
||||
def before_fit(self):
|
||||
if not self.steps_per_epoch: self.steps_per_epoch = len(self.dls.train)
|
||||
self.lrs = [] # store lr values
|
||||
|
||||
self.scheduler = lr_scheduler.OneCycleLR(optimizer = self.opt,
|
||||
max_lr = self.lr_max,
|
||||
total_steps = self.total_steps,
|
||||
epochs = self.n_epochs,
|
||||
steps_per_epoch=self.steps_per_epoch,
|
||||
pct_start=self.pct_start,
|
||||
anneal_strategy=self.anneal_strategy,
|
||||
cycle_momentum=self.cycle_momentum,
|
||||
base_momentum=self.base_momentum,
|
||||
max_momentum=self.max_momentum,
|
||||
div_factor=self.div_factor,
|
||||
final_div_factor=self.final_div_factor,
|
||||
three_phase=self.three_phase,
|
||||
last_epoch=self.last_epoch,
|
||||
verbose=self.verbose
|
||||
)
|
||||
|
||||
def after_batch_train(self):
|
||||
if self.model.training:
|
||||
self.scheduler.step()
|
||||
self.lrs.append( self.scheduler.get_last_lr()[0] )
|
||||
|
||||
def after_fit(self):
|
||||
self.learner.scheduled_lrs = self.lrs
|
||||
|
||||
|
||||
|
||||
class LRFinderCB(Callback):
|
||||
def __init__(self, start_lr=1e-7, end_lr=10, num_iter=100, step_mode='exp', beta=0.98, suggestion='valley'):
|
||||
self.start_lr, self.end_lr = start_lr, end_lr
|
||||
self.num_iter = num_iter
|
||||
self.step_mode = step_mode
|
||||
if beta >= 1: raise ValueError("`num_iter` must be smaller than 1")
|
||||
else: self.beta = beta
|
||||
self.suggestion = suggestion
|
||||
|
||||
def before_fit(self):
|
||||
self.losses, self.lrs = [], []
|
||||
self.best_loss, self.aver_loss = inf, 0
|
||||
self.train_iter = 0
|
||||
|
||||
# save model to load back after fitting
|
||||
self.temp_path = self.save('current', 'temp/', with_opt=False)
|
||||
|
||||
# set base_lr for the optimizer
|
||||
self.set_lr(self.start_lr)
|
||||
|
||||
# check num_iter
|
||||
if not self.num_iter: self.num_iter = len(self.dls.train)
|
||||
# if self.num_iter > len(self.dls.train): self.num_iter = len(self.dls.train)
|
||||
|
||||
# Initialize the proper learning rate policy
|
||||
if self.step_mode.lower() == "exp":
|
||||
self.scheduler = ExponentialLR(self.opt, self.end_lr, self.num_iter)
|
||||
elif self.step_mode.lower() == "linear":
|
||||
self.scheduler = LinearLR(self.opt, self.end_lr, self.num_iter)
|
||||
|
||||
def after_batch_train(self):
|
||||
self.train_iter += 1
|
||||
self.scheduler.step()
|
||||
self.lrs.append( self.scheduler.get_last_lr()[0] )
|
||||
|
||||
# update smooth loss
|
||||
self.smoothing(self.beta)
|
||||
if self.smoothed_loss < self.best_loss: self.best_loss = self.smoothed_loss
|
||||
#Stop if the loss is exploding
|
||||
if self.smoothed_loss > 4 * self.best_loss:
|
||||
raise KeyboardInterrupt # stop fit method
|
||||
if self.train_iter > self.num_iter:
|
||||
raise KeyboardInterrupt # stop fit method
|
||||
|
||||
def smoothing(self, beta):
|
||||
# Smooth the loss if beta is specified
|
||||
self.aver_loss = beta * self.aver_loss + (1-beta) *self.loss.detach().item()
|
||||
self.smoothed_loss = self.aver_loss / (1 - beta**self.train_iter)
|
||||
self.losses.append(self.smoothed_loss)
|
||||
|
||||
def after_fit(self):
|
||||
# reset the gradients
|
||||
self.learner.opt.zero_grad()
|
||||
if self.suggestion == 'valley':
|
||||
self.suggested_lr = valley(self.lrs, self.losses)
|
||||
# load back the model at the previous state
|
||||
self.load(self.temp_path)
|
||||
|
||||
def set_lr(self, lrs):
|
||||
if not isinstance(lrs, list): lrs = [lrs] * len(self.opt.param_groups)
|
||||
if len(lrs) != len(self.opt.param_groups):
|
||||
raise ValueError(
|
||||
"Length of `lrs` is not equal to the number of parameter groups "
|
||||
+ "in the given optimizer")
|
||||
# update lr
|
||||
for param_group, lr in zip(self.opt.param_groups, lrs):
|
||||
param_group["lr"] = lr
|
||||
|
||||
def plot_lr_find(self):
|
||||
import matplotlib.pyplot as plt
|
||||
fig, ax = plt.subplots(1,1)
|
||||
ax.plot(self.lrs, self.losses)
|
||||
ax.set_ylabel("Loss")
|
||||
ax.set_xlabel("Learning Rate")
|
||||
ax.set_xscale('log')
|
||||
plt.grid()
|
||||
|
||||
|
||||
|
||||
class LinearLR(_LRScheduler):
|
||||
"""Linearly increases the learning rate between two boundaries over a number of iterations.
|
||||
|
||||
Arguments:
|
||||
optimizer (torch.optim.Optimizer): wrapped optimizer.
|
||||
end_lr (float): the final learning rate.
|
||||
num_iter (int): the number of iterations over which the test occurs.
|
||||
last_epoch (int, optional): the index of last epoch. Default: -1.
|
||||
"""
|
||||
|
||||
def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
|
||||
self.end_lr = end_lr
|
||||
if num_iter <= 1: raise ValueError("`num_iter` must be larger than 1")
|
||||
self.num_iter = num_iter
|
||||
super(LinearLR, self).__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
r = (self.last_epoch+1) / (self.num_iter - 1)
|
||||
return [base_lr + r * (self.end_lr - base_lr) for base_lr in self.base_lrs]
|
||||
|
||||
|
||||
|
||||
class ExponentialLR(_LRScheduler):
|
||||
"""Exponentially increases the learning rate between two boundaries over a number of iterations.
|
||||
|
||||
Arguments:
|
||||
optimizer (torch.optim.Optimizer): wrapped optimizer.
|
||||
end_lr (float): the final learning rate.
|
||||
num_iter (int): the number of iterations over which the test occurs.
|
||||
last_epoch (int, optional): the index of last epoch. Default: -1.
|
||||
"""
|
||||
|
||||
def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
|
||||
self.end_lr = end_lr
|
||||
self.last_epoch = last_epoch
|
||||
if num_iter <= 1: raise ValueError("`num_iter` must be larger than 1")
|
||||
self.num_iter = num_iter
|
||||
super(ExponentialLR, self).__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
r = (self.last_epoch+1) / (self.num_iter - 1)
|
||||
return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
|
||||
|
||||
|
||||
|
||||
def valley(lrs:list, losses:list):
|
||||
"Suggests a learning rate from the longest valley and returns its index"
|
||||
n = len(losses)
|
||||
max_start, max_end = 0,0
|
||||
|
||||
# find the longest valley
|
||||
lds = [1]*n
|
||||
for i in range(1,n):
|
||||
for j in range(0,i):
|
||||
if (losses[i] < losses[j]) and (lds[i] < lds[j] + 1):
|
||||
lds[i] = lds[j] + 1
|
||||
if lds[max_end] < lds[i]:
|
||||
max_end = i
|
||||
max_start = max_end - lds[max_end]
|
||||
|
||||
sections = (max_end - max_start) / 3
|
||||
idx = max_start + int(sections) + int(sections/2)
|
||||
|
||||
return float(lrs[idx])
|
||||
-284
@@ -1,284 +0,0 @@
|
||||
__all__ = ['TrackTimerCB', 'TrackTrainingCB', 'PrintResultsCB', 'TerminateOnNaNCB',
|
||||
'TrackerCB', 'SaveModelCB', 'EarlyStoppingCB']
|
||||
|
||||
from ..basics import *
|
||||
from .core import Callback
|
||||
import torch
|
||||
import time
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class TrackTimerCB(Callback):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
|
||||
def before_fit(self):
|
||||
self.learner.epoch_time = None
|
||||
|
||||
def before_epoch_train(self):
|
||||
self.start_time = time.time()
|
||||
|
||||
def after_epoch_train(self):
|
||||
self.learner.epoch_time = self.format_time(time.time() - self.start_time)
|
||||
|
||||
def format_time(self, t):
|
||||
"Format `t` (in seconds) to (h):mm:ss"
|
||||
t = int(t)
|
||||
h, m, s = t // 3600, (t // 60) % 60, t % 60
|
||||
if h != 0:
|
||||
return f'{h}:{m:02d}:{s:02d}'
|
||||
else:
|
||||
return f'{m:02d}:{s:02d}'
|
||||
|
||||
|
||||
class TrackTrainingCB(Callback):
|
||||
|
||||
def __init__(self, train_metrics=False, valid_metrics=True):
|
||||
super().__init__()
|
||||
self.train_metrics, self.valid_metrics = train_metrics, valid_metrics
|
||||
|
||||
def init_cb_(self):
|
||||
self.setup()
|
||||
self.initialize_recorder()
|
||||
if hasattr(self.loss_func, 'reduction'):
|
||||
self.mean_reduction_ = True if self.loss_func.reduction == 'mean' else False
|
||||
|
||||
def before_fit(self):
|
||||
self.setup()
|
||||
self.initialize_recorder()
|
||||
if hasattr(self.loss_func, 'reduction'):
|
||||
self.mean_reduction_ = True if self.loss_func.reduction == 'mean' else False
|
||||
|
||||
def setup(self):
|
||||
self.valid_loss = False
|
||||
if self.learner.dls:
|
||||
if not self.learner.dls.valid: self.valid_metrics = False
|
||||
else: self.valid_loss = True
|
||||
|
||||
if self.metrics:
|
||||
if not isinstance(self.metrics, list): self.metrics = [self.metrics]
|
||||
self.metric_names = [func.__name__ for func in self.metrics]
|
||||
else: self.metrics, self.metric_names = [], []
|
||||
|
||||
def initialize_recorder(self):
|
||||
recorder = {'epoch': [], 'train_loss': []}
|
||||
if self.valid_loss: recorder['valid_loss'] = []
|
||||
|
||||
for name in self.metric_names:
|
||||
if self.train_metrics: recorder['train_'+name] = []
|
||||
if self.valid_metrics: recorder['valid_'+name] = []
|
||||
self.recorder = recorder
|
||||
self.learner.recorder = recorder
|
||||
|
||||
|
||||
def initialize_batch_recorder(self, with_metrics):
|
||||
batch_recorder = {'n_samples': [], 'batch_losses': [], 'with_metrics': with_metrics}
|
||||
self.batch_recorder = batch_recorder
|
||||
|
||||
def reset(self):
|
||||
self.targs, self.preds = [],[]
|
||||
self.n_samples = 0
|
||||
self.batch_loss = []
|
||||
|
||||
|
||||
def after_epoch(self):
|
||||
self.recorder['epoch'].append(self.epoch)
|
||||
self.learner.recorder = self.recorder
|
||||
|
||||
def before_epoch_train(self):
|
||||
# define storage for batch training loss and metrics
|
||||
self.initialize_batch_recorder(with_metrics=self.train_metrics)
|
||||
self.reset()
|
||||
|
||||
def before_epoch_valid(self):
|
||||
# if valid data is available, define storage for batch training loss and metrics
|
||||
# if self.dls.valid: self.initialize_batch_recorder(with_metrics=self.valid_metrics)
|
||||
self.initialize_batch_recorder(with_metrics=self.valid_metrics)
|
||||
self.reset()
|
||||
|
||||
|
||||
def after_epoch_train(self):
|
||||
values = self.compute_scores()
|
||||
# save training loss after one epoch
|
||||
self.recorder['train_loss'].append( values['loss'] )
|
||||
# save metrics after one epoch
|
||||
if self.train_metrics:
|
||||
for name, func in zip(self.metric_names, self.metrics):
|
||||
self.recorder['train_'+name].append( values[name] )
|
||||
|
||||
|
||||
def after_epoch_valid(self):
|
||||
# if there is no valid data, don't store
|
||||
if not self.learner.dls.valid: return
|
||||
values = self.compute_scores()
|
||||
# save training loss after one epoch
|
||||
self.recorder['valid_loss'].append( values['loss'] )
|
||||
# save metrics after one epoch
|
||||
if self.valid_metrics:
|
||||
for name, func in zip(self.metric_names, self.metrics):
|
||||
self.recorder['valid_'+name].append( values[name] )
|
||||
|
||||
|
||||
def after_batch_train(self): self.accumulate() # save batch recorder
|
||||
def after_batch_valid(self): self.accumulate()
|
||||
|
||||
def accumulate(self ):
|
||||
xb, yb = self.batch
|
||||
bs = len(xb)
|
||||
self.batch_recorder['n_samples'].append(bs)
|
||||
# get batch loss
|
||||
loss = self.loss.detach()*bs if self.mean_reduction_ else self.loss.detach()
|
||||
self.batch_recorder['batch_losses'].append(loss)
|
||||
|
||||
if yb is None: self.batch_recorder['with_metrics'] = False
|
||||
if len(self.metrics) == 0: self.batch_recorder['with_metrics'] = False
|
||||
# accumulate prediction and target
|
||||
if self.batch_recorder['with_metrics']:
|
||||
self.preds.append(self.pred.detach().cpu())
|
||||
self.targs.append(yb.detach().cpu())
|
||||
|
||||
|
||||
def compute_scores(self):
|
||||
"calculate losses and metrics after each epoch"
|
||||
values = {}
|
||||
# calculate loss after each epoch
|
||||
n = sum(self.batch_recorder['n_samples']) # get total number of samples
|
||||
values['loss'] = sum(self.batch_recorder['batch_losses']).item()/n # averaging
|
||||
|
||||
# calculate metrics if available after each epoch
|
||||
if len(self.preds) == 0: return values
|
||||
self.preds = torch.cat(self.preds)
|
||||
self.targs = torch.cat(self.targs)
|
||||
for func in self.metrics:
|
||||
# values[func.__name__] = func(self.targs, self.preds)
|
||||
values[func.__name__] = func(self.targs, self.preds)
|
||||
return values
|
||||
|
||||
|
||||
class TerminateOnNaNCB(Callback):
|
||||
" A callback to stop the training if loss is NaN"
|
||||
def after_batch_train(self):
|
||||
if torch.isinf(self.loss) or torch.isnan(self.loss): raise KeyboardInterrupt
|
||||
|
||||
|
||||
class PrintResultsCB(Callback):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
def get_header(self, recorder):
|
||||
"recorder is a dictionary"
|
||||
header = list(recorder.keys())
|
||||
return header+['time']
|
||||
|
||||
def before_fit(self):
|
||||
if self.run_finder: return # don't print if lr_finder is called
|
||||
if not hasattr(self.learner, 'recorder'): return # don't print if there is no recorder
|
||||
header = self.get_header(self.learner.recorder)
|
||||
self.print_header = '{:>15s}'*len(header)
|
||||
self.print_value = '{:>15d}' + '{:>15.6f}'*(len(header)-2) + '{:>15}'
|
||||
print(self.print_header.format(*header))
|
||||
|
||||
def after_epoch(self):
|
||||
if self.run_finder: return # don't print if lr_finder is called
|
||||
if not hasattr(self.learner, 'recorder'): return # don't print if there is no recorder
|
||||
epoch_logs = []
|
||||
for key in self.learner.recorder:
|
||||
value=self.learner.recorder[key][-1] if self.learner.recorder[key] else None
|
||||
epoch_logs += [value]
|
||||
if self.learner.epoch_time: epoch_logs.append(self.learner.epoch_time)
|
||||
# print('epoch_logs', epoch_logs)
|
||||
print(self.print_value.format(*epoch_logs))
|
||||
|
||||
|
||||
|
||||
class TrackerCB(Callback):
|
||||
def __init__(self, monitor='train_loss', comp=None, min_delta=0.):
|
||||
super().__init__()
|
||||
if comp is None: comp = np.less if 'loss' in monitor or 'error' in monitor else np.greater
|
||||
if comp == np.less: min_delta *= -1
|
||||
self.monitor, self.comp, self.min_delta = monitor, comp, min_delta
|
||||
|
||||
def before_fit(self):
|
||||
if self.run_finder: return
|
||||
if self.best is None: self.best = float('inf') if self.comp == np.less else -float('inf')
|
||||
self.monitor_names = list(self.learner.recorder.keys())
|
||||
assert self.monitor in self.monitor_names
|
||||
|
||||
def after_epoch(self):
|
||||
if self.run_finder: return
|
||||
val = self.learner.recorder[self.monitor][-1]
|
||||
if self.comp(val - self.min_delta, self.best): self.best, self.new_best = val,True
|
||||
else: self.new_best = False
|
||||
|
||||
|
||||
class SaveModelCB(TrackerCB):
|
||||
def __init__(self, monitor='train_loss', comp=None, min_delta=0.,
|
||||
every_epoch=False, fname='model', path=None, with_opt=False, save_process_id=0, global_rank=None):
|
||||
super().__init__(monitor=monitor, comp=comp, min_delta=min_delta)
|
||||
self.every_epoch = every_epoch
|
||||
self.last_saved_path = None
|
||||
self.path, self.fname = path, fname
|
||||
self.with_opt = with_opt
|
||||
self.save_process_id = save_process_id
|
||||
|
||||
# Identify the worker that saves the model to a file: check if the process' global_rank == save_process_id
|
||||
# If running locally using either a cpu/gpu without using DDP -> set save_process_id = global_rank
|
||||
# Else if running in DDP mode but user doesn't specify global_rank -> global_rank = current_device
|
||||
# (local_rank 0 from each node will save the model)
|
||||
# Else if user provides the global_rank -> use the global_rank to check
|
||||
|
||||
if global_rank:
|
||||
self.global_rank = int(global_rank)
|
||||
else:
|
||||
if torch.cuda.is_available():
|
||||
self.global_rank = torch.cuda.current_device()
|
||||
if not torch.distributed.is_initialized():
|
||||
self.save_process_id = self.global_rank
|
||||
else:
|
||||
self.global_rank = 0
|
||||
|
||||
|
||||
def _save(self, fname, path):
|
||||
if self.global_rank == self.save_process_id:
|
||||
self.last_saved_path = self.learner.save(fname, path, with_opt=self.with_opt)
|
||||
|
||||
def after_epoch(self):
|
||||
if self.every_epoch:
|
||||
if ((self.epoch%self.every_epoch) == 0) or (self.epoch==self.n_epochs-1):
|
||||
self._save(f'{self.fname}_{self.epoch}', self.path)
|
||||
else:
|
||||
super().after_epoch()
|
||||
if self.new_best:
|
||||
print(f'Better model found at epoch {self.epoch} with {self.monitor} value: {self.best}.')
|
||||
self._save(f'{self.fname}', self.path)
|
||||
|
||||
def after_fit(self):
|
||||
if self.run_finder: return
|
||||
if not self.every_epoch and self.global_rank == self.save_process_id:
|
||||
self.learner.load(self.last_saved_path, with_opt=self.with_opt)
|
||||
|
||||
|
||||
class EarlyStoppingCB(TrackerCB):
|
||||
def __init__(self, monitor='train_loss', comp=None, min_delta=0,
|
||||
patient=5):
|
||||
super().__init__(monitor=monitor, comp=comp, min_delta=min_delta)
|
||||
self.patient = patient
|
||||
|
||||
def before_fit(self):
|
||||
# set the impatient level
|
||||
self.impatient_level = 0
|
||||
super().before_fit()
|
||||
|
||||
def after_epoch(self):
|
||||
super().after_epoch()
|
||||
if self.new_best: self.impatient_level = 0 # reset the impatience
|
||||
else:
|
||||
self.impatient_level += 1
|
||||
if self.impatient_level > self.patient:
|
||||
print(f'No improvement since epoch {self.epoch-self.impatient_level}: early stopping')
|
||||
raise KeyboardInterrupt
|
||||
|
||||
|
||||
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .core import Callback
|
||||
from src.models.layers.revin import RevIN
|
||||
|
||||
class RevInCB(Callback):
|
||||
def __init__(self, num_features: int, eps=1e-5,
|
||||
affine:bool=False, denorm:bool=True):
|
||||
"""
|
||||
:param num_features: the number of features or channels
|
||||
:param eps: a value added for numerical stability
|
||||
:param affine: if True, RevIN has learnable affine parameters
|
||||
:param denorm: if True, the output will be de-normalized
|
||||
|
||||
This callback only works with affine=False.
|
||||
if affine=True, the learnable affine_weights and affine_bias are not learnt
|
||||
"""
|
||||
super().__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
self.denorm = denorm
|
||||
self.revin = RevIN(num_features, eps, affine)
|
||||
|
||||
|
||||
def before_forward(self): self.revin_norm()
|
||||
def after_forward(self):
|
||||
if self.denorm: self.revin_denorm()
|
||||
|
||||
def revin_norm(self):
|
||||
xb_revin = self.revin(self.xb, 'norm') # xb_revin: [bs x seq_len x nvars]
|
||||
self.learner.xb = xb_revin
|
||||
|
||||
def revin_denorm(self):
|
||||
pred = self.revin(self.pred, 'denorm') # pred: [bs x target_window x nvars]
|
||||
self.learner.pred = pred
|
||||
|
||||
|
||||
-80
@@ -1,80 +0,0 @@
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
|
||||
class DataLoaders:
|
||||
def __init__(
|
||||
self,
|
||||
datasetCls,
|
||||
dataset_kwargs: dict,
|
||||
batch_size: int,
|
||||
workers: int=0,
|
||||
collate_fn=None,
|
||||
shuffle_train = True,
|
||||
shuffle_val = False
|
||||
):
|
||||
super().__init__()
|
||||
self.datasetCls = datasetCls
|
||||
self.batch_size = batch_size
|
||||
|
||||
if "split" in dataset_kwargs.keys():
|
||||
del dataset_kwargs["split"]
|
||||
self.dataset_kwargs = dataset_kwargs
|
||||
self.workers = workers
|
||||
self.collate_fn = collate_fn
|
||||
self.shuffle_train, self.shuffle_val = shuffle_train, shuffle_val
|
||||
|
||||
self.train = self.train_dataloader()
|
||||
self.valid = self.val_dataloader()
|
||||
self.test = self.test_dataloader()
|
||||
|
||||
|
||||
def train_dataloader(self):
|
||||
return self._make_dloader("train", shuffle=self.shuffle_train)
|
||||
|
||||
def val_dataloader(self):
|
||||
return self._make_dloader("val", shuffle=self.shuffle_val)
|
||||
|
||||
def test_dataloader(self):
|
||||
return self._make_dloader("test", shuffle=False)
|
||||
|
||||
def _make_dloader(self, split, shuffle=False):
|
||||
dataset = self.datasetCls(**self.dataset_kwargs, split=split)
|
||||
if len(dataset) == 0: return None
|
||||
return DataLoader(
|
||||
dataset,
|
||||
shuffle=shuffle,
|
||||
batch_size=self.batch_size,
|
||||
num_workers=self.workers,
|
||||
collate_fn=self.collate_fn,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def add_cli(self, parser):
|
||||
parser.add_argument("--batch_size", type=int, default=128)
|
||||
parser.add_argument(
|
||||
"--workers",
|
||||
type=int,
|
||||
default=6,
|
||||
help="number of parallel workers for pytorch dataloader",
|
||||
)
|
||||
|
||||
def add_dl(self, test_data, batch_size=None, **kwargs):
|
||||
# check of test_data is already a DataLoader
|
||||
from ray.train.torch import _WrappedDataLoader
|
||||
if isinstance(test_data, DataLoader) or isinstance(test_data, _WrappedDataLoader):
|
||||
return test_data
|
||||
|
||||
# get batch_size if not defined
|
||||
if batch_size is None: batch_size=self.batch_size
|
||||
# check if test_data is Dataset, if not, wrap Dataset
|
||||
if not isinstance(test_data, Dataset):
|
||||
test_data = self.train.dataset.new(test_data)
|
||||
|
||||
# create a new DataLoader from Dataset
|
||||
test_data = self.train.new(test_data, batch_size, **kwargs)
|
||||
return test_data
|
||||
|
||||
|
||||
-416
@@ -1,416 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import os
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
|
||||
from src.data.timefeatures import time_features
|
||||
import warnings
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
class Dataset_ETT_hour(Dataset):
|
||||
def __init__(self, root_path, split='train', size=None,
|
||||
features='S', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='h',
|
||||
use_time_features=False
|
||||
):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert split in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[split]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
self.use_time_features = use_time_features
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]
|
||||
border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
|
||||
else: return _torch(seq_x, seq_y)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_ETT_minute(Dataset):
|
||||
def __init__(self, root_path, split='train', size=None,
|
||||
features='S', data_path='ETTm1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='t',
|
||||
use_time_features=False
|
||||
):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert split in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[split]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
self.use_time_features = use_time_features
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]
|
||||
border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
|
||||
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
|
||||
else: return _torch(seq_x, seq_y)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_Custom(Dataset):
|
||||
def __init__(self, root_path, split='train', size=None,
|
||||
features='S', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='h',
|
||||
time_col_name='date', use_time_features=False,
|
||||
train_split=0.7, test_split=0.2
|
||||
):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert split in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[split]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
self.time_col_name = time_col_name
|
||||
self.use_time_features = use_time_features
|
||||
|
||||
# train test ratio
|
||||
self.train_split, self.test_split = train_split, test_split
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
'''
|
||||
df_raw.columns: [time_col_name, ...(other features), target feature]
|
||||
'''
|
||||
cols = list(df_raw.columns)
|
||||
#cols.remove(self.target) if self.target
|
||||
#cols.remove(self.time_col_name)
|
||||
#df_raw = df_raw[[self.time_col_name] + cols + [self.target]]
|
||||
|
||||
num_train = int(len(df_raw) * self.train_split)
|
||||
num_test = int(len(df_raw) * self.test_split)
|
||||
num_vali = len(df_raw) - num_train - num_test
|
||||
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
|
||||
border2s = [num_train, num_train + num_vali, len(df_raw)]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[[self.time_col_name]][border1:border2]
|
||||
df_stamp[self.time_col_name] = pd.to_datetime(df_stamp[self.time_col_name])
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp[self.time_col_name].apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp[self.time_col_name].apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp[self.time_col_name].apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp[self.time_col_name].apply(lambda row: row.hour, 1)
|
||||
data_stamp = df_stamp.drop([self.time_col_name], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp[self.time_col_name].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
if self.use_time_features: return _torch(seq_x, seq_y, seq_x_mark, seq_y_mark)
|
||||
else: return _torch(seq_x, seq_y)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_Pred(Dataset):
|
||||
def __init__(self, root_path, split='pred', size=None,
|
||||
features='S', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, inverse=False, timeenc=0, freq='15min', cols=None):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert split in ['pred']
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.inverse = inverse
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
self.cols = cols
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
'''
|
||||
df_raw.columns: ['date', ...(other features), target feature]
|
||||
'''
|
||||
if self.cols:
|
||||
cols = self.cols.copy()
|
||||
cols.remove(self.target)
|
||||
else:
|
||||
cols = list(df_raw.columns)
|
||||
cols.remove(self.target)
|
||||
cols.remove('date')
|
||||
df_raw = df_raw[['date'] + cols + [self.target]]
|
||||
border1 = len(df_raw) - self.seq_len
|
||||
border2 = len(df_raw)
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
self.scaler.fit(df_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
tmp_stamp = df_raw[['date']][border1:border2]
|
||||
tmp_stamp['date'] = pd.to_datetime(tmp_stamp.date)
|
||||
pred_dates = pd.date_range(tmp_stamp.date.values[-1], periods=self.pred_len + 1, freq=self.freq)
|
||||
|
||||
df_stamp = pd.DataFrame(columns=['date'])
|
||||
df_stamp.date = list(tmp_stamp.date.values) + list(pred_dates[1:])
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
|
||||
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
if self.inverse:
|
||||
self.data_y = df_data.values[border1:border2]
|
||||
else:
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
if self.inverse:
|
||||
seq_y = self.data_x[r_begin:r_begin + self.label_len]
|
||||
else:
|
||||
seq_y = self.data_y[r_begin:r_begin + self.label_len]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
def _torch(*dfs):
|
||||
return tuple(torch.from_numpy(x).float() for x in dfs)
|
||||
-134
@@ -1,134 +0,0 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.tseries import offsets
|
||||
from pandas.tseries.frequencies import to_offset
|
||||
|
||||
|
||||
class TimeFeature:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
pass
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + "()"
|
||||
|
||||
|
||||
class SecondOfMinute(TimeFeature):
|
||||
"""Minute of hour encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.second / 59.0 - 0.5
|
||||
|
||||
|
||||
class MinuteOfHour(TimeFeature):
|
||||
"""Minute of hour encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.minute / 59.0 - 0.5
|
||||
|
||||
|
||||
class HourOfDay(TimeFeature):
|
||||
"""Hour of day encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.hour / 23.0 - 0.5
|
||||
|
||||
|
||||
class DayOfWeek(TimeFeature):
|
||||
"""Hour of day encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.dayofweek / 6.0 - 0.5
|
||||
|
||||
|
||||
class DayOfMonth(TimeFeature):
|
||||
"""Day of month encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.day - 1) / 30.0 - 0.5
|
||||
|
||||
|
||||
class DayOfYear(TimeFeature):
|
||||
"""Day of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.dayofyear - 1) / 365.0 - 0.5
|
||||
|
||||
|
||||
class MonthOfYear(TimeFeature):
|
||||
"""Month of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.month - 1) / 11.0 - 0.5
|
||||
|
||||
|
||||
class WeekOfYear(TimeFeature):
|
||||
"""Week of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.isocalendar().week - 1) / 52.0 - 0.5
|
||||
|
||||
|
||||
def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:
|
||||
"""
|
||||
Returns a list of time features that will be appropriate for the given frequency string.
|
||||
Parameters
|
||||
----------
|
||||
freq_str
|
||||
Frequency string of the form [multiple][granularity] such as "12H", "5min", "1D" etc.
|
||||
"""
|
||||
|
||||
features_by_offsets = {
|
||||
offsets.YearEnd: [],
|
||||
offsets.QuarterEnd: [MonthOfYear],
|
||||
offsets.MonthEnd: [MonthOfYear],
|
||||
offsets.Week: [DayOfMonth, WeekOfYear],
|
||||
offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.Minute: [
|
||||
MinuteOfHour,
|
||||
HourOfDay,
|
||||
DayOfWeek,
|
||||
DayOfMonth,
|
||||
DayOfYear,
|
||||
],
|
||||
offsets.Second: [
|
||||
SecondOfMinute,
|
||||
MinuteOfHour,
|
||||
HourOfDay,
|
||||
DayOfWeek,
|
||||
DayOfMonth,
|
||||
DayOfYear,
|
||||
],
|
||||
}
|
||||
|
||||
offset = to_offset(freq_str)
|
||||
|
||||
for offset_type, feature_classes in features_by_offsets.items():
|
||||
if isinstance(offset, offset_type):
|
||||
return [cls() for cls in feature_classes]
|
||||
|
||||
supported_freq_msg = f"""
|
||||
Unsupported frequency {freq_str}
|
||||
The following frequencies are supported:
|
||||
Y - yearly
|
||||
alias: A
|
||||
M - monthly
|
||||
W - weekly
|
||||
D - daily
|
||||
B - business days
|
||||
H - hourly
|
||||
T - minutely
|
||||
alias: min
|
||||
S - secondly
|
||||
"""
|
||||
raise RuntimeError(supported_freq_msg)
|
||||
|
||||
|
||||
def time_features(dates, freq='h'):
|
||||
return np.vstack([feat(dates) for feat in time_features_from_frequency_str(freq)])
|
||||
-517
@@ -1,517 +0,0 @@
|
||||
|
||||
from typing import List
|
||||
import torch
|
||||
from torch.optim import Adam
|
||||
from torch import nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
|
||||
from .basics import *
|
||||
from .callback.core import *
|
||||
from .callback.tracking import *
|
||||
from .callback.scheduler import *
|
||||
from .callback.distributed import *
|
||||
from .utils import *
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sklearn.base import BaseEstimator
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class Learner(GetAttr):
|
||||
|
||||
def __init__(self, dls, model,
|
||||
loss_func=None,
|
||||
lr=1e-3,
|
||||
cbs=None,
|
||||
metrics=None,
|
||||
opt_func=Adam,
|
||||
**kwargs):
|
||||
|
||||
self.model, self.dls, self.loss_func, self.lr = model, dls, loss_func, lr
|
||||
self.opt_func = opt_func
|
||||
#self.opt = self.opt_func(self.model.parameters(), self.lr)
|
||||
self.set_opt()
|
||||
|
||||
self.metrics = metrics
|
||||
self.n_inp = 2
|
||||
# self.n_inp = self.dls.train.dataset.n_inp if self.dls else 0
|
||||
# Initialize callbacks
|
||||
if cbs and not isinstance(cbs, List): cbs = [cbs]
|
||||
self.initialize_callbacks(cbs)
|
||||
# Indicator of running lr_finder
|
||||
self.run_finder = False
|
||||
|
||||
def set_opt(self):
|
||||
if self.model:
|
||||
self.opt = self.opt_func(self.model.parameters(), self.lr)
|
||||
else: self.opt = None
|
||||
|
||||
|
||||
def default_callback(self):
|
||||
"get a set of default callbacks"
|
||||
default_cbs = [ SetupLearnerCB(), TrackTimerCB(),
|
||||
TrackTrainingCB(train_metrics=False, valid_metrics=True)]
|
||||
return default_cbs
|
||||
|
||||
|
||||
def initialize_callbacks(self, cbs):
|
||||
default_cbs = self.default_callback()
|
||||
self.cbs = update_callbacks(cbs, default_cbs) if cbs else default_cbs
|
||||
# add print CB
|
||||
self.cbs += [PrintResultsCB()]
|
||||
for cb in self.cbs: cb.learner = self
|
||||
self('init_cb')
|
||||
|
||||
|
||||
def add_callback(self, cb):
|
||||
if not cb: return
|
||||
cb.learner = self
|
||||
self.cbs = update_callback(cb, self.cbs)
|
||||
|
||||
def add_callbacks(self, cbs):
|
||||
if not isinstance(cbs, list): cbs = [cbs]
|
||||
for cb in cbs: self.add_callback(cb)
|
||||
|
||||
def remove_callback(self, cb):
|
||||
cb.learn = None
|
||||
self.cbs, removed_cb = remove_callback(cb, self.cbs)
|
||||
return removed_cb
|
||||
|
||||
def remove_callbacks(self, cb_list):
|
||||
for cb in cb_list: self.remove_callback(cb)
|
||||
|
||||
|
||||
def fit(self, n_epochs, lr=None, cbs=None, do_valid=True):
|
||||
" fit the model "
|
||||
self.n_epochs = n_epochs
|
||||
if not self.dls.valid: do_valid = False
|
||||
if cbs: self.add_callbacks(cbs)
|
||||
if lr: self.opt = self.opt_func(self.model.parameters(), lr)
|
||||
|
||||
self('before_fit')
|
||||
try:
|
||||
for self.epoch in range(n_epochs):
|
||||
self('before_epoch')
|
||||
self.one_epoch(train=True)
|
||||
# if self.dls.valid:
|
||||
if do_valid: self.one_epoch(train=False)
|
||||
self('after_epoch')
|
||||
except KeyboardInterrupt: pass
|
||||
self('after_fit')
|
||||
|
||||
|
||||
def fit_one_cycle(self, n_epochs, lr_max=None, pct_start=0.3):
|
||||
self.n_epochs = n_epochs
|
||||
self.lr_max = lr_max if lr_max else self.lr
|
||||
cb = OneCycleLR(lr_max=self.lr_max, pct_start=pct_start)
|
||||
self.fit(self.n_epochs, cbs=cb)
|
||||
|
||||
|
||||
def one_epoch(self, train):
|
||||
self.epoch_train() if train else self.epoch_validate()
|
||||
|
||||
def epoch_train(self):
|
||||
self('before_epoch_train')
|
||||
self.model.train()
|
||||
self.dl = self.dls.train
|
||||
self.all_batches('train')
|
||||
self('after_epoch_train')
|
||||
|
||||
def epoch_validate(self, dl=None):
|
||||
self('before_epoch_valid')
|
||||
# model at evaluation mode
|
||||
self.model.eval()
|
||||
self.dl = dl if dl else self.dls.valid
|
||||
if self.dl:
|
||||
with torch.no_grad(): self.all_batches('valid')
|
||||
self('after_epoch_valid')
|
||||
|
||||
|
||||
def all_batches(self, type_):
|
||||
# for self.num,self.batch in enumerate(progress_bar(dl, leave=False)):
|
||||
for num, batch in enumerate(self.dl):
|
||||
self.iter, self.batch = num, batch
|
||||
if type_ == 'train': self.batch_train()
|
||||
elif type_ == 'valid': self.batch_validate()
|
||||
elif type_ == 'predict': self.batch_predict()
|
||||
elif type_ == 'test': self.batch_test()
|
||||
|
||||
def batch_train(self):
|
||||
self('before_batch_train')
|
||||
self._do_batch_train()
|
||||
self('after_batch_train')
|
||||
|
||||
def batch_validate(self):
|
||||
self('before_batch_valid')
|
||||
self._do_batch_validate()
|
||||
self('after_batch_valid')
|
||||
|
||||
def batch_predict(self):
|
||||
self('before_batch_predict')
|
||||
self._do_batch_predict()
|
||||
self('after_batch_predict')
|
||||
|
||||
def batch_test(self):
|
||||
self('before_batch_test')
|
||||
self._do_batch_test()
|
||||
self('after_batch_test')
|
||||
|
||||
def _do_batch_train(self):
|
||||
# forward + get loss + backward + optimize
|
||||
self.pred, self.loss = self.train_step(self.batch)
|
||||
# zero the parameter gradients
|
||||
self.opt.zero_grad()
|
||||
# gradient
|
||||
self.loss.backward()
|
||||
# update weights
|
||||
self.opt.step()
|
||||
|
||||
def train_step(self, batch):
|
||||
# get the inputs
|
||||
self.xb, self.yb = batch
|
||||
# forward
|
||||
pred = self.model_forward()
|
||||
# compute loss
|
||||
loss = self.loss_func(pred, self.yb)
|
||||
return pred, loss
|
||||
|
||||
def model_forward(self):
|
||||
self('before_forward')
|
||||
self.pred = self.model(self.xb)
|
||||
self('after_forward')
|
||||
return self.pred
|
||||
|
||||
def _do_batch_validate(self):
|
||||
# forward + calculate loss
|
||||
self.pred, self.loss = self.valid_step(self.batch)
|
||||
|
||||
def valid_step(self, batch):
|
||||
# get the inputs
|
||||
self.xb, self.yb = batch
|
||||
# forward
|
||||
pred = self.model_forward()
|
||||
# compute loss
|
||||
loss = self.loss_func(pred, self.yb)
|
||||
return pred, loss
|
||||
|
||||
|
||||
def _do_batch_predict(self):
|
||||
self.pred = self.predict_step(self.batch)
|
||||
|
||||
def predict_step(self, batch):
|
||||
# get the inputs
|
||||
self.xb, self.yb = batch
|
||||
# forward
|
||||
pred = self.model_forward()
|
||||
return pred
|
||||
|
||||
def _do_batch_test(self):
|
||||
self.pred, self.yb = self.test_step(self.batch)
|
||||
|
||||
def test_step(self, batch):
|
||||
# get the inputs
|
||||
self.xb, self.yb = batch
|
||||
# forward
|
||||
pred = self.model_forward()
|
||||
return pred, self.yb
|
||||
|
||||
|
||||
def _predict(self, dl=None):
|
||||
# self('before_validate')
|
||||
self('before_predict')
|
||||
if dl is None: return
|
||||
self.dl = dl
|
||||
self.n_inp = dl.dataset.n_inp
|
||||
self.model.eval() # model at evaluation mode
|
||||
with torch.no_grad(): self.all_batches('predict')
|
||||
self('after_predict')
|
||||
|
||||
|
||||
def predict(self, test_data, weight_path=None, Dataset=None, Dataloader=None, batch_size=None):
|
||||
"""_summary_
|
||||
Args:
|
||||
test_data can be a tensor, numpy array, dataset or dataloader
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
if weight_path is not None: self.load(weight_path)
|
||||
cb = GetPredictionsCB()
|
||||
self.add_callback(cb)
|
||||
test_dl = self._prepare_data(test_data, Dataset, Dataloader, batch_size)
|
||||
self._predict(test_dl)
|
||||
self.preds = cb.preds
|
||||
return to_numpy(self.preds)
|
||||
|
||||
|
||||
def test(self, dl, weight_path=None, scores=None):
|
||||
"""_summary_
|
||||
Args:
|
||||
test_data can be a tensor, numpy array, dataset or dataloader
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
if dl is None: return
|
||||
else: self.dl = dl
|
||||
if weight_path is not None: self.load(weight_path)
|
||||
cb = GetTestCB()
|
||||
self.add_callback(cb)
|
||||
self('before_test')
|
||||
self.model.eval()
|
||||
with torch.no_grad(): self.all_batches('test')
|
||||
self('after_test')
|
||||
self.preds, self.targets = to_numpy([cb.preds, cb.targets])
|
||||
# calculate scores
|
||||
if scores:
|
||||
s_vals = [score(cb.targets, cb.preds).to('cpu').numpy() for score in list(scores)]
|
||||
return self.preds, self.targets, s_vals
|
||||
else: return self.preds, self.targets
|
||||
|
||||
|
||||
def _prepare_data(self, test_data, Dataset=None, Dataloader=None, batch_size=None):
|
||||
if test_data is None: return test_data
|
||||
if Dataset and Dataloader:
|
||||
test_dset = Dataset(test_data)
|
||||
if not batch_size: batch_size=16
|
||||
test_dl = Dataloader(test_dset, batch_size)
|
||||
else:
|
||||
if self.dls:
|
||||
# add test_data to the dataloader defined in the dls.train
|
||||
test_dl = self.dls.add_dl(test_data, batch_size=batch_size)
|
||||
else: test_dl = test_data # assume test_data is already a form of dataloader
|
||||
return test_dl
|
||||
|
||||
|
||||
def get_layer_output(self, inp, layers=None, unwrap=False):
|
||||
"""
|
||||
Args:
|
||||
inp: can be numpy array, torch tensor or dataloader
|
||||
"""
|
||||
self.model.eval()
|
||||
device = next(self.model.parameters()).device
|
||||
if isinstance(inp, np.ndarray): inp = torch.Tensor(inp).to(device)
|
||||
if isinstance(inp, torch.Tensor): inp = inp.to(device)
|
||||
|
||||
return get_layer_output(inp, model=self.model, layers=layers, unwrap=unwrap)
|
||||
|
||||
|
||||
def fine_tune(self, n_epochs, base_lr=None, freeze_epochs=1, pct_start=0.3):
|
||||
"""
|
||||
fintune the pretrained model. First the entire model is freezed, only head is trained
|
||||
up to a freeze_epochs number. Then the model is unfreezed and the entire model is trained
|
||||
"""
|
||||
assert (n_epochs>0)|(freeze_epochs>0), "Either n_epochs or freeze_epochs has to be > 0"
|
||||
if not base_lr: base_lr = self.lr
|
||||
# Finetune the head of freeze_epochs > 0:
|
||||
if freeze_epochs > 0:
|
||||
print('Finetune the head')
|
||||
self.freeze()
|
||||
self.fit_one_cycle(freeze_epochs, lr_max=base_lr, pct_start=pct_start)
|
||||
|
||||
# Finetune the entire network if n_epochs > 0
|
||||
if n_epochs > 0:
|
||||
print('Finetune the entire network')
|
||||
self.unfreeze()
|
||||
self.fit_one_cycle(n_epochs, lr_max=base_lr/2, pct_start=pct_start)
|
||||
|
||||
|
||||
def linear_probe(self, n_epochs, base_lr=None, pct_start=0.3):
|
||||
"""
|
||||
linear probing the pretrained model. The model is freeze except the head during finetuning
|
||||
"""
|
||||
assert (n_epochs>0), "n_epochs has to be > 0"
|
||||
if not base_lr: base_lr = self.lr
|
||||
print('Finetune the head')
|
||||
self.freeze()
|
||||
self.fit_one_cycle(n_epochs, lr_max=base_lr, pct_start=pct_start)
|
||||
|
||||
|
||||
def lr_finder(self, start_lr=1e-7, end_lr=10, num_iter=100, step_mode='exp', show_plot=True, suggestion='valley'):
|
||||
"""
|
||||
find the learning rate
|
||||
"""
|
||||
n_epochs = num_iter//len(self.dls.train) + 1
|
||||
# indicator of lr_finder method is applied
|
||||
self.run_finder = True
|
||||
# add LRFinderCB to callback list and will remove later
|
||||
cb = LRFinderCB(start_lr, end_lr, num_iter, step_mode, suggestion=suggestion)
|
||||
# fit
|
||||
self.fit(n_epochs=n_epochs, cbs=cb, do_valid=False)
|
||||
# should remove LRFinderCB callback after fitting
|
||||
self.remove_callback(cb)
|
||||
self.run_finder = False
|
||||
if show_plot: cb.plot_lr_find()
|
||||
if suggestion: return cb.suggested_lr
|
||||
|
||||
|
||||
|
||||
def freeze(self):
|
||||
"""
|
||||
freeze the model head
|
||||
require the model to have head attribute
|
||||
"""
|
||||
if hasattr(get_model(self.model), 'head'):
|
||||
# print('model head is available')
|
||||
for param in get_model(self.model).parameters(): param.requires_grad = False
|
||||
for param in get_model(self.model).head.parameters(): param.requires_grad = True
|
||||
# print('model is frozen except the head')
|
||||
|
||||
|
||||
def unfreeze(self):
|
||||
for param in get_model(self.model).parameters(): param.requires_grad = True
|
||||
|
||||
|
||||
def __call__(self, name):
|
||||
for cb in self.cbs:
|
||||
attr = getattr(cb, name)
|
||||
if attr is not None: attr()
|
||||
|
||||
|
||||
def save(self, fname, path, **kwargs):
|
||||
"""
|
||||
Save model and optimizer state (if `with_opt`) to `self.path/file`
|
||||
"""
|
||||
fname = join_path_file(fname, path, ext='.pth')
|
||||
save_model(fname, self.model, getattr(self,'opt',None), **kwargs)
|
||||
return fname
|
||||
|
||||
|
||||
def load(self, fname, with_opt=False, device='cuda', strict=True, **kwargs):
|
||||
"""
|
||||
load the model
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
device = "cpu"
|
||||
load_model(fname, self.model, self.opt, with_opt, device=device, strict=strict)
|
||||
|
||||
|
||||
def get_params(self, deep=True, **kwargs):
|
||||
params = BaseEstimator.get_params(self, deep=deep, **kwargs)
|
||||
return params
|
||||
|
||||
def _get_param_names(self):
|
||||
return (k for k in self.__dict__ if not k.endswith('_'))
|
||||
|
||||
|
||||
def set_params(self, **kwargs):
|
||||
params = {}
|
||||
for key, val in kwargs.items():
|
||||
params[key] = val
|
||||
BaseEstimator.set_params(self, **params)
|
||||
|
||||
def to_distributed(self,
|
||||
sync_bn=True, # Whether to replace all batch norm with `nn.SyncBatchNorm`
|
||||
**kwargs
|
||||
):
|
||||
local_rank = int(os.environ.get('LOCAL_RANK'))
|
||||
world_size = int(os.environ.get('WORLD_SIZE'))
|
||||
rank = int(os.environ.get('RANK'))
|
||||
print('Process {} (out of {})'.format(
|
||||
rank, torch.distributed.get_world_size()))
|
||||
|
||||
self.add_callback(DistributedTrainer(local_rank=local_rank, world_size=world_size, sync_bn=sync_bn, **kwargs))
|
||||
|
||||
return self
|
||||
|
||||
|
||||
def save_model(path, model, opt, with_opt=True, pickle_protocol=2):
|
||||
"Save `model` to `file` along with `opt` (if available, and if `with_opt`)"
|
||||
if opt is None: with_opt=False
|
||||
state = get_model(model).state_dict()
|
||||
if with_opt: state = {'model': state, 'opt':opt.state_dict()}
|
||||
torch.save(state, path, pickle_protocol=pickle_protocol)
|
||||
|
||||
|
||||
def load_model(path, model, opt=None, with_opt=False, device='cpu', strict=True):
|
||||
" load the saved model "
|
||||
state = torch.load(path, map_location=device)
|
||||
if not opt: with_opt=False
|
||||
model_state = state['model'] if with_opt else state
|
||||
get_model(model).load_state_dict(model_state, strict=strict)
|
||||
if with_opt: opt.load_state_dict(state['opt'])
|
||||
model = model.to(device)
|
||||
|
||||
|
||||
def join_path_file(file, path, ext=''):
|
||||
"Return `path/file` if file is a string or a `Path`, file otherwise"
|
||||
if not isinstance(file, (str, Path)): return file
|
||||
if not isinstance(path, Path): path = Path(path)
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path/f'{file}{ext}'
|
||||
|
||||
|
||||
def get_model(model):
|
||||
"Return the model maybe wrapped inside `model`."
|
||||
return model.module if isinstance(model, (DistributedDataParallel, nn.DataParallel)) else model
|
||||
|
||||
|
||||
def transfer_weights(weights_path, model, exclude_head=True, device='cpu'):
|
||||
# state_dict = model.state_dict()
|
||||
new_state_dict = torch.load(weights_path, map_location=device)
|
||||
#print('new_state_dict',new_state_dict)
|
||||
matched_layers = 0
|
||||
unmatched_layers = []
|
||||
for name, param in model.state_dict().items():
|
||||
if exclude_head and 'head' in name: continue
|
||||
if name in new_state_dict:
|
||||
matched_layers += 1
|
||||
input_param = new_state_dict[name]
|
||||
if input_param.shape == param.shape: param.copy_(input_param)
|
||||
else: unmatched_layers.append(name)
|
||||
else:
|
||||
unmatched_layers.append(name)
|
||||
pass # these are weights that weren't in the original model, such as a new head
|
||||
if matched_layers == 0: raise Exception("No shared weight names were found between the models")
|
||||
else:
|
||||
if len(unmatched_layers) > 0:
|
||||
print(f'check unmatched_layers: {unmatched_layers}')
|
||||
else:
|
||||
print(f"weights from {weights_path} successfully transferred!\n")
|
||||
model = model.to(device)
|
||||
return model
|
||||
|
||||
|
||||
def update_callback(cb, list_cbs):
|
||||
for cb_ in list_cbs:
|
||||
if type(cb_) == type(cb): list_cbs.remove(cb_)
|
||||
list_cbs += [cb]
|
||||
return list_cbs
|
||||
|
||||
def update_callbacks(list_cbs, default_cbs):
|
||||
for cb in list_cbs: default_cbs = update_callback(cb, default_cbs)
|
||||
return default_cbs
|
||||
|
||||
def remove_callback(cb, list_cbs):
|
||||
for cb_ in list_cbs:
|
||||
if type(cb_) == type(cb):
|
||||
list_cbs.remove(cb_)
|
||||
break
|
||||
return list_cbs, cb_
|
||||
|
||||
|
||||
def get_layer_output(inp, model, layers=None, unwrap=False):
|
||||
"""
|
||||
layers is a list of module names
|
||||
"""
|
||||
orig_model = model
|
||||
|
||||
if unwrap: model = unwrap_model(model)
|
||||
if not layers: layers = list(dict(model.named_children()).keys())
|
||||
if not isinstance(layers, list): layers = [layers]
|
||||
|
||||
activation = {}
|
||||
def getActivation(name):
|
||||
# the hook signature
|
||||
def hook(model, input, output):
|
||||
activation[name] = output.detach().cpu().numpy()
|
||||
return hook
|
||||
|
||||
# register forward hooks on the layers of choice
|
||||
h_list = [getattr(model, layer).register_forward_hook(getActivation(layer)) for layer in layers]
|
||||
|
||||
model.eval()
|
||||
out = orig_model(inp)
|
||||
for h in h_list: h.remove()
|
||||
return activation
|
||||
-21
@@ -1,21 +0,0 @@
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
|
||||
def mse(y_true, y_pred):
|
||||
return F.mse_loss(y_true, y_pred, reduction='mean')
|
||||
|
||||
def rmse(y_true, y_pred):
|
||||
return torch.sqrt(F.mse_loss(y_true, y_pred, reduction='mean'))
|
||||
|
||||
def mae(y_true, y_pred):
|
||||
return F.l1_loss(y_true, y_pred, reduction='mean')
|
||||
|
||||
def r2_score(y_true, y_pred):
|
||||
from sklearn.metrics import r2_score
|
||||
return r2_score(y_true, y_pred)
|
||||
|
||||
def mape(y_true, y_pred):
|
||||
from sklearn.metrics import mean_absolute_percentage_error
|
||||
return mean_absolute_percentage_error(y_true, y_pred)
|
||||
-1
@@ -1 +0,0 @@
|
||||
|
||||
Vendored
-115
@@ -1,115 +0,0 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
from typing import Callable, Optional
|
||||
|
||||
class MultiheadAttention(nn.Module):
|
||||
def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0., qkv_bias=True, lsa=False):
|
||||
"""Multi Head Attention Layer
|
||||
Input shape:
|
||||
Q: [batch_size (bs) x max_q_len x d_model]
|
||||
K, V: [batch_size (bs) x q_len x d_model]
|
||||
mask: [q_len x q_len]
|
||||
"""
|
||||
super().__init__()
|
||||
d_k = d_model // n_heads if d_k is None else d_k
|
||||
d_v = d_model // n_heads if d_v is None else d_v
|
||||
|
||||
self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
|
||||
|
||||
self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
|
||||
self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
|
||||
self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)
|
||||
|
||||
# Scaled Dot-Product Attention (multiple heads)
|
||||
self.res_attention = res_attention
|
||||
self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout, res_attention=self.res_attention, lsa=lsa)
|
||||
|
||||
# Poject output
|
||||
self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))
|
||||
|
||||
|
||||
def forward(self, Q:Tensor, K:Optional[Tensor]=None, V:Optional[Tensor]=None, prev:Optional[Tensor]=None,
|
||||
key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):
|
||||
|
||||
bs = Q.size(0)
|
||||
if K is None: K = Q
|
||||
if V is None: V = Q
|
||||
|
||||
# Linear (+ split in multiple heads)
|
||||
q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,2) # q_s : [bs x n_heads x max_q_len x d_k]
|
||||
k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0,2,3,1) # k_s : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
|
||||
v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1,2) # v_s : [bs x n_heads x q_len x d_v]
|
||||
|
||||
# Apply Scaled Dot-Product Attention (multiple heads)
|
||||
if self.res_attention:
|
||||
output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
else:
|
||||
output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
|
||||
# output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
|
||||
|
||||
# back to the original inputs dimensions
|
||||
output = output.transpose(1, 2).contiguous().view(bs, -1, self.n_heads * self.d_v) # output: [bs x q_len x n_heads * d_v]
|
||||
output = self.to_out(output)
|
||||
|
||||
if self.res_attention: return output, attn_weights, attn_scores
|
||||
else: return output, attn_weights
|
||||
|
||||
|
||||
|
||||
class ScaledDotProductAttention(nn.Module):
|
||||
r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
|
||||
(Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
|
||||
by Lee et al, 2021)"""
|
||||
|
||||
def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
|
||||
super().__init__()
|
||||
self.attn_dropout = nn.Dropout(attn_dropout)
|
||||
self.res_attention = res_attention
|
||||
head_dim = d_model // n_heads
|
||||
self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
|
||||
self.lsa = lsa
|
||||
|
||||
def forward(self, q:Tensor, k:Tensor, v:Tensor, prev:Optional[Tensor]=None, key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):
|
||||
'''
|
||||
Input shape:
|
||||
q : [bs x n_heads x max_q_len x d_k]
|
||||
k : [bs x n_heads x d_k x seq_len]
|
||||
v : [bs x n_heads x seq_len x d_v]
|
||||
prev : [bs x n_heads x q_len x seq_len]
|
||||
key_padding_mask: [bs x seq_len]
|
||||
attn_mask : [1 x seq_len x seq_len]
|
||||
Output shape:
|
||||
output: [bs x n_heads x q_len x d_v]
|
||||
attn : [bs x n_heads x q_len x seq_len]
|
||||
scores : [bs x n_heads x q_len x seq_len]
|
||||
'''
|
||||
|
||||
# Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
|
||||
attn_scores = torch.matmul(q, k) * self.scale # attn_scores : [bs x n_heads x max_q_len x q_len]
|
||||
|
||||
# Add pre-softmax attention scores from the previous layer (optional)
|
||||
if prev is not None: attn_scores = attn_scores + prev
|
||||
|
||||
# Attention mask (optional)
|
||||
if attn_mask is not None: # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_scores.masked_fill_(attn_mask, -np.inf)
|
||||
else:
|
||||
attn_scores += attn_mask
|
||||
|
||||
# Key padding mask (optional)
|
||||
if key_padding_mask is not None: # mask with shape [bs x q_len] (only when max_w_len == q_len)
|
||||
attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)
|
||||
|
||||
# normalize the attention weights
|
||||
attn_weights = F.softmax(attn_scores, dim=-1) # attn_weights : [bs x n_heads x max_q_len x q_len]
|
||||
attn_weights = self.attn_dropout(attn_weights)
|
||||
|
||||
# compute the new values given the attention weights
|
||||
output = torch.matmul(attn_weights, v) # output: [bs x n_heads x max_q_len x d_v]
|
||||
|
||||
if self.res_attention: return output, attn_weights, attn_scores
|
||||
else: return output, attn_weights
|
||||
|
||||
-47
@@ -1,47 +0,0 @@
|
||||
__all__ = ['Transpose', 'LinBnDrop', 'SigmoidRange', 'sigmoid_range', 'get_activation_fn']
|
||||
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class Transpose(nn.Module):
|
||||
def __init__(self, *dims, contiguous=False):
|
||||
super().__init__()
|
||||
self.dims, self.contiguous = dims, contiguous
|
||||
def forward(self, x):
|
||||
if self.contiguous: return x.transpose(*self.dims).contiguous()
|
||||
else: return x.transpose(*self.dims)
|
||||
|
||||
|
||||
class SigmoidRange(nn.Module):
|
||||
def __init__(self, low, high):
|
||||
super().__init__()
|
||||
self.low, self.high = low, high
|
||||
# self.low, self.high = ranges
|
||||
def forward(self, x):
|
||||
# return sigmoid_range(x, self.low, self.high)
|
||||
return torch.sigmoid(x) * (self.high - self.low) + self.low
|
||||
|
||||
|
||||
class LinBnDrop(nn.Sequential):
|
||||
"Module grouping `BatchNorm1d`, `Dropout` and `Linear` layers"
|
||||
def __init__(self, n_in, n_out, bn=True, p=0., act=None, lin_first=False):
|
||||
layers = [nn.BatchNorm2d(n_out if lin_first else n_in, ndim=1)] if bn else []
|
||||
if p != 0: layers.append(nn.Dropout(p))
|
||||
lin = [nn.Linear(n_in, n_out, bias=not bn)]
|
||||
if act is not None: lin.append(act)
|
||||
layers = lin+layers if lin_first else layers+lin
|
||||
super().__init__(*layers)
|
||||
|
||||
|
||||
def sigmoid_range(x, low, high):
|
||||
"Sigmoid function with range `(low, high)`"
|
||||
return torch.sigmoid(x) * (high - low) + low
|
||||
|
||||
def get_activation_fn(activation):
|
||||
if callable(activation): return activation()
|
||||
elif activation.lower() == "relu": return nn.ReLU()
|
||||
elif activation.lower() == "gelu": return nn.GELU()
|
||||
raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable')
|
||||
|
||||
|
||||
-103
@@ -1,103 +0,0 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class LinearRegressionHead(nn.Module):
|
||||
def __init__(self, n_vars, d_model, output_dim, head_dropout, y_range=None):
|
||||
super().__init__()
|
||||
self.y_range = y_range
|
||||
self.flatten = nn.Flatten(start_dim=1)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
self.linear = nn.Linear(n_vars*d_model, output_dim)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x output_dim]
|
||||
"""
|
||||
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
|
||||
x = self.flatten(x) # x: bs x nvars * d_model
|
||||
x = self.dropout(x)
|
||||
y = self.linear(x) # y: bs x output_dim
|
||||
if self.y_range: y = SigmoidRange(*self.y_range)(y)
|
||||
return y
|
||||
|
||||
|
||||
class LinearClassificationHead(nn.Module):
|
||||
def __init__(self, n_vars, d_model, n_classes, head_dropout):
|
||||
super().__init__()
|
||||
self.flatten = nn.Flatten(start_dim=1)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
self.linear = nn.Linear(n_vars*d_model, n_classes)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x n_classes]
|
||||
"""
|
||||
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
|
||||
x = self.flatten(x) # x: bs x nvars * d_model
|
||||
x = self.dropout(x)
|
||||
y = self.linear(x) # y: bs x n_classes
|
||||
return y
|
||||
|
||||
|
||||
class LinearPredictionHead(nn.Module):
|
||||
def __init__(self, individual, n_vars, d_model, num_patch, forecast_len, head_dropout=0):
|
||||
super().__init__()
|
||||
|
||||
self.individual = individual
|
||||
self.n_vars = n_vars
|
||||
head_dim = d_model*num_patch
|
||||
|
||||
if self.individual:
|
||||
self.linears = nn.ModuleList()
|
||||
self.dropouts = nn.ModuleList()
|
||||
self.flattens = nn.ModuleList()
|
||||
for i in range(self.n_vars):
|
||||
self.flattens.append(nn.Flatten(start_dim=-2))
|
||||
self.linears.append(nn.Linear(head_dim, forecast_len))
|
||||
self.dropouts.append(nn.Dropout(head_dropout))
|
||||
else:
|
||||
self.flatten = nn.Flatten(start_dim=-2)
|
||||
self.linear = nn.Linear(head_dim, forecast_len)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x forecast_len x nvars]
|
||||
"""
|
||||
if self.individual:
|
||||
x_out = []
|
||||
for i in range(self.n_vars):
|
||||
z = self.flattens[i](x[:,i,:,:]) # z: [bs x d_model * num_patch]
|
||||
z = self.linears[i](z) # z: [bs x forecast_len]
|
||||
z = self.dropouts[i](z)
|
||||
x_out.append(z)
|
||||
x = torch.stack(x_out, dim=1) # x: [bs x nvars x forecast_len]
|
||||
else:
|
||||
x = self.flatten(x)
|
||||
x = self.dropout(x)
|
||||
x = self.linear(x)
|
||||
return x.transpose(2,1) # [bs x forecast_len x nvars]
|
||||
|
||||
|
||||
class LinearPretrainHead(nn.Module):
|
||||
def __init__(self, d_model, patch_len, dropout):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear = nn.Linear(d_model, patch_len)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: tensor [bs x nvars x d_model x num_patch]
|
||||
output: tensor [bs x nvars x num_patch x patch_len]
|
||||
"""
|
||||
|
||||
x = x.transpose(2,3) # [bs x nvars x num_patch x d_model]
|
||||
x = self.linear( self.dropout(x) ) # [bs x nvars x num_patch x patch_len]
|
||||
x = x.permute(0,2,1,3) # [bs x num_patch x nvars x patch_len]
|
||||
return x
|
||||
|
||||
Vendored
-46
@@ -1,46 +0,0 @@
|
||||
__all__ = ['PositionalEncoding', 'SinCosPosEncoding', 'positional_encoding']
|
||||
|
||||
# Cell
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import math
|
||||
|
||||
# Cell
|
||||
def PositionalEncoding(q_len, d_model, normalize=True):
|
||||
pe = torch.zeros(q_len, d_model)
|
||||
position = torch.arange(0, q_len).unsqueeze(1)
|
||||
div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
|
||||
pe[:, 0::2] = torch.sin(position * div_term)
|
||||
pe[:, 1::2] = torch.cos(position * div_term)
|
||||
if normalize:
|
||||
pe = pe - pe.mean()
|
||||
pe = pe / (pe.std() * 10)
|
||||
return pe
|
||||
|
||||
SinCosPosEncoding = PositionalEncoding
|
||||
|
||||
|
||||
def positional_encoding(pe, learn_pe, q_len, d_model):
|
||||
# Positional encoding
|
||||
if pe == None:
|
||||
W_pos = torch.empty((q_len, d_model)) # pe = None and learn_pe = False can be used to measure impact of pe
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
learn_pe = False
|
||||
elif pe == 'zero':
|
||||
W_pos = torch.empty((q_len, 1))
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
elif pe == 'zeros':
|
||||
W_pos = torch.empty((q_len, d_model))
|
||||
nn.init.uniform_(W_pos, -0.02, 0.02)
|
||||
elif pe == 'normal' or pe == 'gauss':
|
||||
W_pos = torch.zeros((q_len, 1))
|
||||
torch.nn.init.normal_(W_pos, mean=0.0, std=0.1)
|
||||
elif pe == 'uniform':
|
||||
W_pos = torch.zeros((q_len, 1))
|
||||
nn.init.uniform_(W_pos, a=0.0, b=0.1)
|
||||
elif pe == 'sincos': W_pos = PositionalEncoding(q_len, d_model, normalize=True)
|
||||
else: raise ValueError(f"{pe} is not a valid pe (positional encoder. Available types: 'gauss'=='normal', \
|
||||
'zeros', 'zero', uniform', 'sincos', None.)")
|
||||
return nn.Parameter(W_pos, requires_grad=learn_pe)
|
||||
|
||||
-51
@@ -1,51 +0,0 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
class RevIN(nn.Module):
|
||||
def __init__(self, num_features: int, eps=1e-5, affine=True):
|
||||
"""
|
||||
:param num_features: the number of features or channels
|
||||
:param eps: a value added for numerical stability
|
||||
:param affine: if True, RevIN has learnable affine parameters
|
||||
"""
|
||||
super(RevIN, self).__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
if self.affine:
|
||||
self._init_params()
|
||||
|
||||
def forward(self, x, mode:str):
|
||||
if mode == 'norm':
|
||||
self._get_statistics(x)
|
||||
x = self._normalize(x)
|
||||
elif mode == 'denorm':
|
||||
x = self._denormalize(x)
|
||||
else: raise NotImplementedError
|
||||
return x
|
||||
|
||||
def _init_params(self):
|
||||
# initialize RevIN params: (C,)
|
||||
self.affine_weight = nn.Parameter(torch.ones(self.num_features))
|
||||
self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
|
||||
|
||||
def _get_statistics(self, x):
|
||||
dim2reduce = tuple(range(1, x.ndim-1))
|
||||
self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
|
||||
self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
|
||||
|
||||
def _normalize(self, x):
|
||||
x = x - self.mean
|
||||
x = x / self.stdev
|
||||
if self.affine:
|
||||
x = x * self.affine_weight
|
||||
x = x + self.affine_bias
|
||||
return x
|
||||
|
||||
def _denormalize(self, x):
|
||||
if self.affine:
|
||||
x = x - self.affine_bias
|
||||
x = x / (self.affine_weight + self.eps*self.eps)
|
||||
x = x * self.stdev
|
||||
x = x + self.mean
|
||||
return x
|
||||
-333
@@ -1,333 +0,0 @@
|
||||
|
||||
__all__ = ['PatchTST']
|
||||
|
||||
# Cell
|
||||
from typing import Callable, Optional
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
from collections import OrderedDict
|
||||
from ..models.layers.pos_encoding import *
|
||||
from ..models.layers.basics import *
|
||||
from ..models.layers.attention import *
|
||||
|
||||
|
||||
# Cell
|
||||
class PatchTST(nn.Module):
|
||||
"""
|
||||
Output dimension:
|
||||
[bs x target_dim x nvars] for prediction
|
||||
[bs x target_dim] for regression
|
||||
[bs x target_dim] for classification
|
||||
[bs x num_patch x n_vars x patch_len] for pretrain
|
||||
"""
|
||||
def __init__(self, c_in:int, target_dim:int, patch_len:int, stride:int, num_patch:int,
|
||||
n_layers:int=3, d_model=128, n_heads=16, shared_embedding=True, d_ff:int=256,
|
||||
norm:str='BatchNorm', attn_dropout:float=0., dropout:float=0., act:str="gelu",
|
||||
res_attention:bool=True, pre_norm:bool=False, store_attn:bool=False,
|
||||
pe:str='zeros', learn_pe:bool=True, head_dropout = 0,
|
||||
head_type = "prediction", individual = False,
|
||||
y_range:Optional[tuple]=None, verbose:bool=False, **kwargs):
|
||||
|
||||
super().__init__()
|
||||
|
||||
assert head_type in ['pretrain', 'prediction', 'regression', 'classification'], 'head type should be either pretrain, prediction, or regression'
|
||||
# Backbone
|
||||
self.backbone = PatchTSTEncoder(c_in, num_patch=num_patch, patch_len=patch_len,
|
||||
n_layers=n_layers, d_model=d_model, n_heads=n_heads,
|
||||
shared_embedding=shared_embedding, d_ff=d_ff,
|
||||
attn_dropout=attn_dropout, dropout=dropout, act=act,
|
||||
res_attention=res_attention, pre_norm=pre_norm, store_attn=store_attn,
|
||||
pe=pe, learn_pe=learn_pe, verbose=verbose, **kwargs)
|
||||
|
||||
# Head
|
||||
self.n_vars = c_in
|
||||
self.head_type = head_type
|
||||
|
||||
if head_type == "pretrain":
|
||||
self.head = PretrainHead(d_model, patch_len, head_dropout) # custom head passed as a partial func with all its kwargs
|
||||
elif head_type == "prediction":
|
||||
self.head = PredictionHead(individual, self.n_vars, d_model, num_patch, target_dim, head_dropout)
|
||||
elif head_type == "regression":
|
||||
self.head = RegressionHead(self.n_vars, d_model, target_dim, head_dropout, y_range)
|
||||
elif head_type == "classification":
|
||||
self.head = ClassificationHead(self.n_vars, d_model, target_dim, head_dropout)
|
||||
|
||||
|
||||
def forward(self, z):
|
||||
"""
|
||||
z: tensor [bs x num_patch x n_vars x patch_len]
|
||||
"""
|
||||
z = self.backbone(z) # z: [bs x nvars x d_model x num_patch]
|
||||
z = self.head(z)
|
||||
# z: [bs x target_dim x nvars] for prediction
|
||||
# [bs x target_dim] for regression
|
||||
# [bs x target_dim] for classification
|
||||
# [bs x num_patch x n_vars x patch_len] for pretrain
|
||||
return z
|
||||
|
||||
|
||||
class RegressionHead(nn.Module):
|
||||
def __init__(self, n_vars, d_model, output_dim, head_dropout, y_range=None):
|
||||
super().__init__()
|
||||
self.y_range = y_range
|
||||
self.flatten = nn.Flatten(start_dim=1)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
self.linear = nn.Linear(n_vars*d_model, output_dim)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x output_dim]
|
||||
"""
|
||||
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
|
||||
x = self.flatten(x) # x: bs x nvars * d_model
|
||||
x = self.dropout(x)
|
||||
y = self.linear(x) # y: bs x output_dim
|
||||
if self.y_range: y = SigmoidRange(*self.y_range)(y)
|
||||
return y
|
||||
|
||||
|
||||
class ClassificationHead(nn.Module):
|
||||
def __init__(self, n_vars, d_model, n_classes, head_dropout):
|
||||
super().__init__()
|
||||
self.flatten = nn.Flatten(start_dim=1)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
self.linear = nn.Linear(n_vars*d_model, n_classes)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x n_classes]
|
||||
"""
|
||||
x = x[:,:,:,-1] # only consider the last item in the sequence, x: bs x nvars x d_model
|
||||
x = self.flatten(x) # x: bs x nvars * d_model
|
||||
x = self.dropout(x)
|
||||
y = self.linear(x) # y: bs x n_classes
|
||||
return y
|
||||
|
||||
|
||||
class PredictionHead(nn.Module):
|
||||
def __init__(self, individual, n_vars, d_model, num_patch, forecast_len, head_dropout=0, flatten=False):
|
||||
super().__init__()
|
||||
|
||||
self.individual = individual
|
||||
self.n_vars = n_vars
|
||||
self.flatten = flatten
|
||||
head_dim = d_model*num_patch
|
||||
|
||||
if self.individual:
|
||||
self.linears = nn.ModuleList()
|
||||
self.dropouts = nn.ModuleList()
|
||||
self.flattens = nn.ModuleList()
|
||||
for i in range(self.n_vars):
|
||||
self.flattens.append(nn.Flatten(start_dim=-2))
|
||||
self.linears.append(nn.Linear(head_dim, forecast_len))
|
||||
self.dropouts.append(nn.Dropout(head_dropout))
|
||||
else:
|
||||
self.flatten = nn.Flatten(start_dim=-2)
|
||||
self.linear = nn.Linear(head_dim, forecast_len)
|
||||
self.dropout = nn.Dropout(head_dropout)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: [bs x nvars x d_model x num_patch]
|
||||
output: [bs x forecast_len x nvars]
|
||||
"""
|
||||
if self.individual:
|
||||
x_out = []
|
||||
for i in range(self.n_vars):
|
||||
z = self.flattens[i](x[:,i,:,:]) # z: [bs x d_model * num_patch]
|
||||
z = self.linears[i](z) # z: [bs x forecast_len]
|
||||
z = self.dropouts[i](z)
|
||||
x_out.append(z)
|
||||
x = torch.stack(x_out, dim=1) # x: [bs x nvars x forecast_len]
|
||||
else:
|
||||
x = self.flatten(x) # x: [bs x nvars x (d_model * num_patch)]
|
||||
x = self.dropout(x)
|
||||
x = self.linear(x) # x: [bs x nvars x forecast_len]
|
||||
return x.transpose(2,1) # [bs x forecast_len x nvars]
|
||||
|
||||
|
||||
class PretrainHead(nn.Module):
|
||||
def __init__(self, d_model, patch_len, dropout):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear = nn.Linear(d_model, patch_len)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
x: tensor [bs x nvars x d_model x num_patch]
|
||||
output: tensor [bs x nvars x num_patch x patch_len]
|
||||
"""
|
||||
|
||||
x = x.transpose(2,3) # [bs x nvars x num_patch x d_model]
|
||||
x = self.linear( self.dropout(x) ) # [bs x nvars x num_patch x patch_len]
|
||||
x = x.permute(0,2,1,3) # [bs x num_patch x nvars x patch_len]
|
||||
return x
|
||||
|
||||
|
||||
class PatchTSTEncoder(nn.Module):
|
||||
def __init__(self, c_in, num_patch, patch_len,
|
||||
n_layers=3, d_model=128, n_heads=16, shared_embedding=True,
|
||||
d_ff=256, norm='BatchNorm', attn_dropout=0., dropout=0., act="gelu", store_attn=False,
|
||||
res_attention=True, pre_norm=False,
|
||||
pe='zeros', learn_pe=True, verbose=False, **kwargs):
|
||||
|
||||
super().__init__()
|
||||
self.n_vars = c_in
|
||||
self.num_patch = num_patch
|
||||
self.patch_len = patch_len
|
||||
self.d_model = d_model
|
||||
self.shared_embedding = shared_embedding
|
||||
|
||||
# Input encoding: projection of feature vectors onto a d-dim vector space
|
||||
if not shared_embedding:
|
||||
self.W_P = nn.ModuleList()
|
||||
for _ in range(self.n_vars): self.W_P.append(nn.Linear(patch_len, d_model))
|
||||
else:
|
||||
self.W_P = nn.Linear(patch_len, d_model)
|
||||
|
||||
# Positional encoding
|
||||
self.W_pos = positional_encoding(pe, learn_pe, num_patch, d_model)
|
||||
|
||||
# Residual dropout
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
# Encoder
|
||||
self.encoder = TSTEncoder(d_model, n_heads, d_ff=d_ff, norm=norm, attn_dropout=attn_dropout, dropout=dropout,
|
||||
pre_norm=pre_norm, activation=act, res_attention=res_attention, n_layers=n_layers,
|
||||
store_attn=store_attn)
|
||||
|
||||
def forward(self, x) -> Tensor:
|
||||
"""
|
||||
x: tensor [bs x num_patch x nvars x patch_len]
|
||||
"""
|
||||
bs, num_patch, n_vars, patch_len = x.shape
|
||||
# Input encoding
|
||||
if not self.shared_embedding:
|
||||
x_out = []
|
||||
for i in range(n_vars):
|
||||
z = self.W_P[i](x[:,:,i,:])
|
||||
x_out.append(z)
|
||||
x = torch.stack(x_out, dim=2)
|
||||
else:
|
||||
x = self.W_P(x) # x: [bs x num_patch x nvars x d_model]
|
||||
x = x.transpose(1,2) # x: [bs x nvars x num_patch x d_model]
|
||||
|
||||
u = torch.reshape(x, (bs*n_vars, num_patch, self.d_model) ) # u: [bs * nvars x num_patch x d_model]
|
||||
u = self.dropout(u + self.W_pos) # u: [bs * nvars x num_patch x d_model]
|
||||
|
||||
# Encoder
|
||||
z = self.encoder(u) # z: [bs * nvars x num_patch x d_model]
|
||||
z = torch.reshape(z, (-1,n_vars, num_patch, self.d_model)) # z: [bs x nvars x num_patch x d_model]
|
||||
z = z.permute(0,1,3,2) # z: [bs x nvars x d_model x num_patch]
|
||||
|
||||
return z
|
||||
|
||||
|
||||
# Cell
|
||||
class TSTEncoder(nn.Module):
|
||||
def __init__(self, d_model, n_heads, d_ff=None,
|
||||
norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu',
|
||||
res_attention=False, n_layers=1, pre_norm=False, store_attn=False):
|
||||
super().__init__()
|
||||
|
||||
self.layers = nn.ModuleList([TSTEncoderLayer(d_model, n_heads=n_heads, d_ff=d_ff, norm=norm,
|
||||
attn_dropout=attn_dropout, dropout=dropout,
|
||||
activation=activation, res_attention=res_attention,
|
||||
pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)])
|
||||
self.res_attention = res_attention
|
||||
|
||||
def forward(self, src:Tensor):
|
||||
"""
|
||||
src: tensor [bs x q_len x d_model]
|
||||
"""
|
||||
output = src
|
||||
scores = None
|
||||
if self.res_attention:
|
||||
for mod in self.layers: output, scores = mod(output, prev=scores)
|
||||
return output
|
||||
else:
|
||||
for mod in self.layers: output = mod(output)
|
||||
return output
|
||||
|
||||
|
||||
|
||||
class TSTEncoderLayer(nn.Module):
|
||||
def __init__(self, d_model, n_heads, d_ff=256, store_attn=False,
|
||||
norm='BatchNorm', attn_dropout=0, dropout=0., bias=True,
|
||||
activation="gelu", res_attention=False, pre_norm=False):
|
||||
super().__init__()
|
||||
assert not d_model%n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
|
||||
d_k = d_model // n_heads
|
||||
d_v = d_model // n_heads
|
||||
|
||||
# Multi-Head attention
|
||||
self.res_attention = res_attention
|
||||
self.self_attn = MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout, res_attention=res_attention)
|
||||
|
||||
# Add & Norm
|
||||
self.dropout_attn = nn.Dropout(dropout)
|
||||
if "batch" in norm.lower():
|
||||
self.norm_attn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
|
||||
else:
|
||||
self.norm_attn = nn.LayerNorm(d_model)
|
||||
|
||||
# Position-wise Feed-Forward
|
||||
self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
|
||||
get_activation_fn(activation),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(d_ff, d_model, bias=bias))
|
||||
|
||||
# Add & Norm
|
||||
self.dropout_ffn = nn.Dropout(dropout)
|
||||
if "batch" in norm.lower():
|
||||
self.norm_ffn = nn.Sequential(Transpose(1,2), nn.BatchNorm1d(d_model), Transpose(1,2))
|
||||
else:
|
||||
self.norm_ffn = nn.LayerNorm(d_model)
|
||||
|
||||
self.pre_norm = pre_norm
|
||||
self.store_attn = store_attn
|
||||
|
||||
|
||||
def forward(self, src:Tensor, prev:Optional[Tensor]=None):
|
||||
"""
|
||||
src: tensor [bs x q_len x d_model]
|
||||
"""
|
||||
# Multi-Head attention sublayer
|
||||
if self.pre_norm:
|
||||
src = self.norm_attn(src)
|
||||
## Multi-Head attention
|
||||
if self.res_attention:
|
||||
src2, attn, scores = self.self_attn(src, src, src, prev)
|
||||
else:
|
||||
src2, attn = self.self_attn(src, src, src)
|
||||
if self.store_attn:
|
||||
self.attn = attn
|
||||
## Add & Norm
|
||||
src = src + self.dropout_attn(src2) # Add: residual connection with residual dropout
|
||||
if not self.pre_norm:
|
||||
src = self.norm_attn(src)
|
||||
|
||||
# Feed-forward sublayer
|
||||
if self.pre_norm:
|
||||
src = self.norm_ffn(src)
|
||||
## Position-wise Feed-Forward
|
||||
src2 = self.ff(src)
|
||||
## Add & Norm
|
||||
src = src + self.dropout_ffn(src2) # Add: residual connection with residual dropout
|
||||
if not self.pre_norm:
|
||||
src = self.norm_ffn(src)
|
||||
|
||||
if self.res_attention:
|
||||
return src, scores
|
||||
else:
|
||||
return src
|
||||
|
||||
|
||||
|
||||
@@ -1,58 +0,0 @@
|
||||
|
||||
from torch import nn
|
||||
import collections
|
||||
from collections import OrderedDict
|
||||
import torch
|
||||
import os
|
||||
from datetime import timedelta
|
||||
|
||||
|
||||
def init_ddp():
|
||||
local_rank = int(os.environ.get('LOCAL_RANK'))
|
||||
world_size = int(os.environ.get('WORLD_SIZE'))
|
||||
rank = int(os.environ.get('RANK'))
|
||||
|
||||
torch.cuda.set_device(local_rank)
|
||||
torch.distributed.init_process_group(
|
||||
'nccl',
|
||||
init_method='env://',
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
timeout=timedelta(seconds=600)
|
||||
)
|
||||
|
||||
|
||||
def nested_children(m: nn.Module):
|
||||
children = dict(m.named_children())
|
||||
output = {}
|
||||
if children == {}:
|
||||
# if module has no children; m is last child! :O
|
||||
return m
|
||||
else:
|
||||
# look for children from children... to the last child!
|
||||
for name, child in children.items():
|
||||
try:
|
||||
output[name] = nested_children(child)
|
||||
except TypeError:
|
||||
output[name] = nested_children(child)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def flatten_dict(d, parent_key='', sep='_'):
|
||||
items = []
|
||||
for k, v in d.items():
|
||||
new_key = parent_key + sep + k if parent_key else k
|
||||
if isinstance(v, collections.MutableMapping):
|
||||
items.extend(flatten_dict(v, new_key, sep=sep).items())
|
||||
else:
|
||||
items.append((new_key, v))
|
||||
return dict(items)
|
||||
|
||||
|
||||
def unwrap_model(model):
|
||||
unwrapped_model = nested_children(model)
|
||||
unwrapped_model = flatten_dict(unwrapped_model)
|
||||
unwrapped_model = nn.Sequential(OrderedDict(unwrapped_model))
|
||||
return unwrapped_model
|
||||
|
||||
-209
@@ -1,209 +0,0 @@
|
||||
# ALL scripts in this file come from Autoformer
|
||||
if [ ! -d "./logs" ]; then
|
||||
mkdir ./logs
|
||||
fi
|
||||
|
||||
if [ ! -d "./logs/LongForecasting" ]; then
|
||||
mkdir ./logs/LongForecasting
|
||||
fi
|
||||
|
||||
random_seed=2021
|
||||
model_name=Autoformer
|
||||
|
||||
for pred_len in 96 192 336 720
|
||||
do
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path exchange_rate.csv \
|
||||
--model_id exchange_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 8 \
|
||||
--dec_in 8 \
|
||||
--c_out 8 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 1 >logs/LongForecasting/$model_name'_exchange_rate_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path electricity.csv \
|
||||
--model_id electricity_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 321 \
|
||||
--dec_in 321 \
|
||||
--c_out 321 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_electricity_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path traffic.csv \
|
||||
--model_id traffic_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 862 \
|
||||
--dec_in 862 \
|
||||
--c_out 862 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 3 >logs/LongForecasting/$model_name'_traffic_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path weather.csv \
|
||||
--model_id weather_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 21 \
|
||||
--dec_in 21 \
|
||||
--c_out 21 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 2 >logs/LongForecasting/$model_name'_weather_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTh1.csv \
|
||||
--model_id ETTh1_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTh1 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Etth1_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTh2.csv \
|
||||
--model_id ETTh2_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTh2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Etth2_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTm1.csv \
|
||||
--model_id ETTm1_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTm1 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Ettm1_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTm2.csv \
|
||||
--model_id ETTm2_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTm2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Ettm2_'$pred_len.log
|
||||
done
|
||||
|
||||
for pred_len in 24 36 48 60
|
||||
do
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path national_illness.csv \
|
||||
--model_id ili_36_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 36 \
|
||||
--label_len 18 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_ili_'$pred_len.log
|
||||
done
|
||||
-21
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2021 xxxx
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
-53
@@ -1,53 +0,0 @@
|
||||
# FEDformer
|
||||
From https://github.com/MAZiqing/FEDformer
|
||||
|
||||
|
||||
Frequency Enhanced Decomposed
|
||||
Transformer (FEDformer) is more efficient than
|
||||
standard Transformer with a linear complexity
|
||||
to the sequence length.
|
||||
|
||||
Our empirical studies
|
||||
with six benchmark datasets show that compared
|
||||
with state-of-the-art methods, FEDformer can
|
||||
reduce prediction error by 14.8% and 22.6%
|
||||
for multivariate and univariate time series,
|
||||
respectively.
|
||||
|
||||
|
||||
## Get Started
|
||||
|
||||
1. Install Python 3.6, PyTorch 1.9.0.
|
||||
2. Download data. You can obtain all the six benchmarks from xxxx.
|
||||
3. Train the model. We provide the experiment scripts of all benchmarks under the folder `./scripts`. You can reproduce the experiment results by:
|
||||
|
||||
```bash
|
||||
bash ./scripts/run_M.sh
|
||||
bash ./scripts/run_S.sh
|
||||
```
|
||||
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repo useful, please cite our paper.
|
||||
|
||||
```
|
||||
xxxxx
|
||||
```
|
||||
|
||||
## Contact
|
||||
|
||||
If you have any question or want to use the code, please contact xxx@xxxx .
|
||||
|
||||
## Acknowledgement
|
||||
|
||||
We appreciate the following github repos a lot for their valuable code base or datasets:
|
||||
|
||||
https://github.com/thuml/Autoformer
|
||||
|
||||
https://github.com/zhouhaoyi/Informer2020
|
||||
|
||||
https://github.com/zhouhaoyi/ETDataset
|
||||
|
||||
https://github.com/laiguokun/multivariate-time-series-data
|
||||
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
from data_provider.data_loader import Dataset_ETT_hour, Dataset_ETT_minute, Dataset_Custom,Dataset_sin
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
data_dict = {
|
||||
'ETTh1': Dataset_ETT_hour,
|
||||
'ETTh2': Dataset_ETT_hour,
|
||||
'ETTm1': Dataset_ETT_minute,
|
||||
'ETTm2': Dataset_ETT_minute,
|
||||
'custom': Dataset_Custom,
|
||||
'sin':Dataset_sin,
|
||||
}
|
||||
|
||||
|
||||
def data_provider(args, flag):
|
||||
Data = data_dict[args.data]
|
||||
timeenc = 0 if args.embed != 'timeF' else 1
|
||||
|
||||
if flag == 'test':
|
||||
shuffle_flag = False
|
||||
drop_last = True
|
||||
batch_size = args.batch_size
|
||||
freq = args.freq
|
||||
elif flag == 'pred':
|
||||
shuffle_flag = False
|
||||
drop_last = False
|
||||
batch_size = 1
|
||||
freq = args.detail_freq
|
||||
Data = Dataset_Pred
|
||||
else:
|
||||
shuffle_flag = True
|
||||
drop_last = True
|
||||
batch_size = args.batch_size
|
||||
freq = args.freq
|
||||
|
||||
data_set = Data(
|
||||
root_path=args.root_path,
|
||||
data_path=args.data_path,
|
||||
flag=flag,
|
||||
size=[args.seq_len, args.label_len, args.pred_len],
|
||||
features=args.features,
|
||||
target=args.target,
|
||||
timeenc=timeenc,
|
||||
freq=freq
|
||||
)
|
||||
print(flag, len(data_set))
|
||||
data_loader = DataLoader(
|
||||
data_set,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle_flag,
|
||||
num_workers=args.num_workers,
|
||||
drop_last=drop_last)
|
||||
return data_set, data_loader
|
||||
-372
@@ -1,372 +0,0 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import os
|
||||
import torch
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
from utils.timefeatures import time_features
|
||||
import warnings
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
class Dataset_ETT_hour(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None,
|
||||
features='S', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='h'):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len]
|
||||
border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_ETT_minute(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None,
|
||||
features='S', data_path='ETTm1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='t'):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len]
|
||||
border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
df_stamp['minute'] = df_stamp.date.apply(lambda row: row.minute, 1)
|
||||
df_stamp['minute'] = df_stamp.minute.map(lambda x: x // 15)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_Custom(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None,
|
||||
features='S', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, timeenc=0, freq='h'):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
'''
|
||||
df_raw.columns: ['date', ...(other features), target feature]
|
||||
'''
|
||||
cols = list(df_raw.columns)
|
||||
cols.remove(self.target)
|
||||
cols.remove('date')
|
||||
df_raw = df_raw[['date'] + cols + [self.target]]
|
||||
# print(cols)
|
||||
num_train = int(len(df_raw) * 0.7)
|
||||
num_test = int(len(df_raw) * 0.2)
|
||||
num_vali = len(df_raw) - num_train - num_test
|
||||
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
|
||||
border2s = [num_train, num_train + num_vali, len(df_raw)]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
class Dataset_sin(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None,
|
||||
features='S', data_path='sin.csv',
|
||||
target='y', scale=True, timeenc=0, freq='h'):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.label_len = 24 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.label_len = size[1]
|
||||
self.pred_len = size[2]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
'''
|
||||
df_raw.columns: ['date', ...(other features), target feature]
|
||||
'''
|
||||
cols = list(df_raw.columns)
|
||||
print(cols)
|
||||
cols.remove(self.target)
|
||||
cols.remove('x')
|
||||
df_raw = df_raw[['x'] + cols + [self.target]]
|
||||
# print(cols)
|
||||
num_train = int(len(df_raw) * 0.7)
|
||||
num_test = int(len(df_raw) * 0.2)
|
||||
num_vali = len(df_raw) - num_train - num_test
|
||||
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
|
||||
border2s = [num_train, num_train + num_vali, len(df_raw)]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end - self.label_len
|
||||
r_end = r_begin + self.label_len + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = np.zeros_like(seq_x)
|
||||
seq_y_mark = np.zeros_like(seq_y)
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
Vendored
-37
@@ -1,37 +0,0 @@
|
||||
import os
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Exp_Basic(object):
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
self.device = self._acquire_device()
|
||||
self.model = self._build_model().to(self.device)
|
||||
|
||||
def _build_model(self):
|
||||
raise NotImplementedError
|
||||
return None
|
||||
|
||||
def _acquire_device(self):
|
||||
if self.args.use_gpu:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(
|
||||
self.args.gpu) if not self.args.use_multi_gpu else self.args.devices
|
||||
device = torch.device('cuda:{}'.format(self.args.gpu))
|
||||
print('Use GPU: cuda:{}'.format(self.args.gpu))
|
||||
else:
|
||||
device = torch.device('cpu')
|
||||
print('Use CPU')
|
||||
return device
|
||||
|
||||
def _get_data(self):
|
||||
pass
|
||||
|
||||
def vali(self):
|
||||
pass
|
||||
|
||||
def train(self):
|
||||
pass
|
||||
|
||||
def test(self):
|
||||
pass
|
||||
Vendored
-347
@@ -1,347 +0,0 @@
|
||||
from data_provider.data_factory import data_provider
|
||||
from exp.exp_basic import Exp_Basic
|
||||
from models import FEDformer, Autoformer, Informer, Transformer
|
||||
from utils.tools import EarlyStopping, adjust_learning_rate, visual
|
||||
from utils.metrics import metric
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import optim
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import warnings
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import io
|
||||
from scipy import stats
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
class Exp_Main(Exp_Basic):
|
||||
def __init__(self, args):
|
||||
super(Exp_Main, self).__init__(args)
|
||||
|
||||
def _build_model(self):
|
||||
model_dict = {
|
||||
'FEDformer': FEDformer,
|
||||
'Autoformer': Autoformer,
|
||||
'Transformer': Transformer,
|
||||
'Informer': Informer,
|
||||
}
|
||||
model = model_dict[self.args.model].Model(self.args).float()
|
||||
|
||||
if self.args.use_multi_gpu and self.args.use_gpu:
|
||||
model = nn.DataParallel(model, device_ids=self.args.device_ids)
|
||||
return model
|
||||
|
||||
def _get_data(self, flag):
|
||||
data_set, data_loader = data_provider(self.args, flag)
|
||||
return data_set, data_loader
|
||||
|
||||
def _select_optimizer(self):
|
||||
model_optim = optim.Adam(self.model.parameters(), lr=self.args.learning_rate)
|
||||
return model_optim
|
||||
|
||||
def _select_criterion(self):
|
||||
criterion = nn.MSELoss()
|
||||
return criterion
|
||||
|
||||
def vali(self, vali_data, vali_loader, criterion):
|
||||
total_loss = []
|
||||
ks_test_96,ks_test_192,ks_test_336,ks_test_720,ks_test_96_back=[],[],[],[],[]
|
||||
ks_result=[]
|
||||
ks_test_96_raw,ks_test_192_raw,ks_test_336_raw,ks_test_720_raw,ks_test_96_back_raw=[],[],[],[],[]
|
||||
self.model.eval()
|
||||
input_len=720
|
||||
with torch.no_grad():
|
||||
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(vali_loader):
|
||||
batch_x = batch_x.float().to(self.device)
|
||||
batch_y = batch_y.float()
|
||||
|
||||
batch_x_mark = batch_x_mark.float().to(self.device)
|
||||
batch_y_mark = batch_y_mark.float().to(self.device)
|
||||
|
||||
# decoder input
|
||||
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
|
||||
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
|
||||
# encoder - decoder
|
||||
if self.args.use_amp:
|
||||
with torch.cuda.amp.autocast():
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
else:
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
f_dim = -1 if self.args.features == 'MS' else 0
|
||||
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
|
||||
|
||||
pred = outputs.detach().cpu()
|
||||
true = batch_y.detach().cpu()
|
||||
|
||||
loss = criterion(pred, true)
|
||||
|
||||
total_loss.append(loss)
|
||||
total_loss = np.average(total_loss)
|
||||
self.model.train()
|
||||
return total_loss
|
||||
|
||||
def train(self, setting):
|
||||
train_data, train_loader = self._get_data(flag='train')
|
||||
vali_data, vali_loader = self._get_data(flag='val')
|
||||
test_data, test_loader = self._get_data(flag='test')
|
||||
|
||||
path = os.path.join(self.args.checkpoints, setting)
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
|
||||
time_now = time.time()
|
||||
|
||||
train_steps = len(train_loader)
|
||||
early_stopping = EarlyStopping(patience=self.args.patience, verbose=True)
|
||||
|
||||
model_optim = self._select_optimizer()
|
||||
criterion = self._select_criterion()
|
||||
|
||||
if self.args.use_amp:
|
||||
scaler = torch.cuda.amp.GradScaler()
|
||||
|
||||
for epoch in range(self.args.train_epochs):
|
||||
iter_count = 0
|
||||
train_loss = []
|
||||
|
||||
self.model.train()
|
||||
epoch_time = time.time()
|
||||
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(train_loader):
|
||||
iter_count += 1
|
||||
model_optim.zero_grad()
|
||||
batch_x = batch_x.float().to(self.device)
|
||||
|
||||
batch_y = batch_y.float().to(self.device)
|
||||
batch_x_mark = batch_x_mark.float().to(self.device)
|
||||
batch_y_mark = batch_y_mark.float().to(self.device)
|
||||
|
||||
# decoder input
|
||||
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
|
||||
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
|
||||
|
||||
# encoder - decoder
|
||||
if self.args.use_amp:
|
||||
with torch.cuda.amp.autocast():
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
|
||||
f_dim = -1 if self.args.features == 'MS' else 0
|
||||
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
|
||||
loss = criterion(outputs, batch_y)
|
||||
train_loss.append(loss.item())
|
||||
else:
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
|
||||
f_dim = -1 if self.args.features == 'MS' else 0
|
||||
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
|
||||
|
||||
# if i==0:
|
||||
# from scipy import stats
|
||||
# pred = outputs.detach().cpu().numpy()
|
||||
# true = batch_y.detach().cpu().numpy()
|
||||
# input_data = batch_x.detach().cpu().numpy()
|
||||
# plot_index1=np.arange(input_data.shape[1])
|
||||
# plot_index2=np.arange(input_data.shape[1],input_data.shape[1]+pred.shape[1])
|
||||
# plt.cla()
|
||||
# plt.plot(plot_index1,input_data[0,:,-1:],label='input')
|
||||
# plt.plot(plot_index2,pred[0,:,-1:],label="predict")
|
||||
# plt.plot(plot_index2,true[0,:,-1:],label="true")
|
||||
# print('KS test1',stats.kstest(input_data[0,:,-1:].reshape(-1),pred[0,-input_data.shape[1]:,-1:].reshape(-1)))
|
||||
|
||||
# plt.legend()
|
||||
# #f = io.BytesIO()
|
||||
# plt.savefig("sample0.png",format="png")
|
||||
# #plt.clf()
|
||||
# plt.cla()
|
||||
# plt.plot(plot_index1,input_data[8,:,-1:],label='input')
|
||||
# plt.plot(plot_index2,pred[8,:,-1:],label="predict")
|
||||
# plt.plot(plot_index2,true[8,:,-1:],label="true")
|
||||
# plt.legend()
|
||||
# plt.savefig("sample1.png",format="png")
|
||||
# print('KS test2',stats.kstest(input_data[8,:,-1:].reshape(-1),pred[8,-input_data.shape[1]:,-1:].reshape(-1)))
|
||||
# raise Exception('aaa')
|
||||
|
||||
|
||||
loss = criterion(outputs, batch_y)
|
||||
train_loss.append(loss.item())
|
||||
|
||||
if (i + 1) % 100 == 0:
|
||||
# print("\titers: {0}, epoch: {1} | loss: {2:.7f}".format(i + 1, epoch + 1, loss.item()))
|
||||
speed = (time.time() - time_now) / iter_count
|
||||
left_time = speed * ((self.args.train_epochs - epoch) * train_steps - i)
|
||||
# print('\tspeed: {:.4f}s/iter; left time: {:.4f}s'.format(speed, left_time))
|
||||
iter_count = 0
|
||||
time_now = time.time()
|
||||
|
||||
if self.args.use_amp:
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(model_optim)
|
||||
scaler.update()
|
||||
else:
|
||||
loss.backward()
|
||||
model_optim.step()
|
||||
|
||||
print("Epoch: {} cost time: {}".format(epoch + 1, time.time() - epoch_time))
|
||||
train_loss = np.average(train_loss)
|
||||
vali_loss = self.vali(vali_data, vali_loader, criterion)
|
||||
test_loss = self.vali(test_data, test_loader, criterion)
|
||||
|
||||
print("Epoch: {0}, Steps: {1} | Train Loss: {2:.7f} Vali Loss: {3:.7f} Test Loss: {4:.7f}".format(
|
||||
epoch + 1, train_steps, train_loss, vali_loss, test_loss))
|
||||
early_stopping(vali_loss, self.model, path)
|
||||
if early_stopping.early_stop:
|
||||
print("Early stopping")
|
||||
break
|
||||
|
||||
adjust_learning_rate(model_optim, epoch + 1, self.args)
|
||||
|
||||
best_model_path = path + '/' + 'checkpoint.pth'
|
||||
self.model.load_state_dict(torch.load(best_model_path))
|
||||
|
||||
return self.model
|
||||
|
||||
def test(self, setting, test=0):
|
||||
test_data, test_loader = self._get_data(flag='test')
|
||||
if test:
|
||||
print('loading model')
|
||||
self.model.load_state_dict(torch.load(os.path.join('./checkpoints/' + setting, 'checkpoint.pth')))
|
||||
|
||||
preds = []
|
||||
trues = []
|
||||
folder_path = './test_results/' + setting + '/'
|
||||
if not os.path.exists(folder_path):
|
||||
os.makedirs(folder_path)
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(test_loader):
|
||||
batch_x = batch_x.float().to(self.device)
|
||||
batch_y = batch_y.float().to(self.device)
|
||||
|
||||
batch_x_mark = batch_x_mark.float().to(self.device)
|
||||
batch_y_mark = batch_y_mark.float().to(self.device)
|
||||
|
||||
# decoder input
|
||||
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
|
||||
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
|
||||
# encoder - decoder
|
||||
if self.args.use_amp:
|
||||
with torch.cuda.amp.autocast():
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
else:
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
|
||||
f_dim = -1 if self.args.features == 'MS' else 0
|
||||
|
||||
batch_y = batch_y[:, -self.args.pred_len:, f_dim:].to(self.device)
|
||||
outputs = outputs.detach().cpu().numpy()
|
||||
batch_y = batch_y.detach().cpu().numpy()
|
||||
|
||||
pred = outputs # outputs.detach().cpu().numpy() # .squeeze()
|
||||
true = batch_y # batch_y.detach().cpu().numpy() # .squeeze()
|
||||
|
||||
preds.append(pred)
|
||||
trues.append(true)
|
||||
if i % 20 == 0:
|
||||
input = batch_x.detach().cpu().numpy()
|
||||
gt = np.concatenate((input[0, :, -1], true[0, :, -1]), axis=0)
|
||||
pd = np.concatenate((input[0, :, -1], pred[0, :, -1]), axis=0)
|
||||
visual(gt, pd, os.path.join(folder_path, str(i) + '.pdf'))
|
||||
|
||||
preds = np.array(preds)
|
||||
trues = np.array(trues)
|
||||
print('test shape:', preds.shape, trues.shape)
|
||||
preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
|
||||
trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])
|
||||
print('test shape:', preds.shape, trues.shape)
|
||||
|
||||
# result save
|
||||
folder_path = './results/' + setting + '/'
|
||||
if not os.path.exists(folder_path):
|
||||
os.makedirs(folder_path)
|
||||
|
||||
mae, mse, rmse, mape, mspe = metric(preds, trues)
|
||||
print('mse:{}, mae:{}'.format(mse, mae))
|
||||
f = open("result.txt", 'a')
|
||||
f.write(setting + " \n")
|
||||
f.write('mse:{}, mae:{}'.format(mse, mae))
|
||||
f.write('\n')
|
||||
f.write('\n')
|
||||
f.close()
|
||||
|
||||
np.save(folder_path + 'metrics.npy', np.array([mae, mse, rmse, mape, mspe]))
|
||||
np.save(folder_path + 'pred.npy', preds)
|
||||
np.save(folder_path + 'true.npy', trues)
|
||||
|
||||
return
|
||||
|
||||
def predict(self, setting, load=False):
|
||||
pred_data, pred_loader = self._get_data(flag='pred')
|
||||
|
||||
if load:
|
||||
path = os.path.join(self.args.checkpoints, setting)
|
||||
best_model_path = path + '/' + 'checkpoint.pth'
|
||||
self.model.load_state_dict(torch.load(best_model_path))
|
||||
|
||||
preds = []
|
||||
|
||||
self.model.eval()
|
||||
with torch.no_grad():
|
||||
for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(pred_loader):
|
||||
batch_x = batch_x.float().to(self.device)
|
||||
batch_y = batch_y.float()
|
||||
batch_x_mark = batch_x_mark.float().to(self.device)
|
||||
batch_y_mark = batch_y_mark.float().to(self.device)
|
||||
|
||||
# decoder input
|
||||
dec_inp = torch.zeros_like(batch_y[:, -self.args.pred_len:, :]).float()
|
||||
dec_inp = torch.cat([batch_y[:, :self.args.label_len, :], dec_inp], dim=1).float().to(self.device)
|
||||
# encoder - decoder
|
||||
|
||||
if self.args.use_amp:
|
||||
with torch.cuda.amp.autocast():
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
else:
|
||||
if self.args.output_attention:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
|
||||
else:
|
||||
outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
|
||||
pred = outputs.detach().cpu().numpy() # .squeeze()
|
||||
preds.append(pred)
|
||||
|
||||
preds = np.array(preds)
|
||||
preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
|
||||
|
||||
# result save
|
||||
folder_path = './results/' + setting + '/'
|
||||
if not os.path.exists(folder_path):
|
||||
os.makedirs(folder_path)
|
||||
|
||||
np.save(folder_path + 'real_prediction.npy', preds)
|
||||
|
||||
return
|
||||
-222
@@ -1,222 +0,0 @@
|
||||
import time
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
import math
|
||||
from torch.nn.functional import interpolate
|
||||
|
||||
|
||||
def decor_time(func):
|
||||
def func2(*args, **kw):
|
||||
now = time.time()
|
||||
y = func(*args, **kw)
|
||||
t = time.time() - now
|
||||
print('call <{}>, time={}'.format(func.__name__, t))
|
||||
return y
|
||||
return func2
|
||||
|
||||
|
||||
class AutoCorrelation(nn.Module):
|
||||
"""
|
||||
AutoCorrelation Mechanism with the following two phases:
|
||||
(1) period-based dependencies discovery
|
||||
(2) time delay aggregation
|
||||
This block can replace the self-attention family mechanism seamlessly.
|
||||
"""
|
||||
def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False, configs=None):
|
||||
super(AutoCorrelation, self).__init__()
|
||||
print('Autocorrelation used !')
|
||||
self.factor = factor
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
self.agg = None
|
||||
self.use_wavelet = configs.wavelet
|
||||
|
||||
# @decor_time
|
||||
def time_delay_agg_training(self, values, corr):
|
||||
"""
|
||||
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
||||
This is for the training phase.
|
||||
"""
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
||||
index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
|
||||
weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
pattern = torch.roll(tmp_values, -int(index[i]), -1)
|
||||
delays_agg = delays_agg + pattern * \
|
||||
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
||||
return delays_agg # size=[B, H, d, S]
|
||||
|
||||
def time_delay_agg_inference(self, values, corr):
|
||||
"""
|
||||
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
||||
This is for the inference phase.
|
||||
"""
|
||||
batch = values.shape[0]
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# index init
|
||||
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
||||
weights = torch.topk(mean_value, top_k, dim=-1)[0]
|
||||
delay = torch.topk(mean_value, top_k, dim=-1)[1]
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values.repeat(1, 1, 1, 2)
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
|
||||
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
||||
delays_agg = delays_agg + pattern * \
|
||||
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
||||
return delays_agg
|
||||
|
||||
def time_delay_agg_full(self, values, corr):
|
||||
"""
|
||||
Standard version of Autocorrelation
|
||||
"""
|
||||
batch = values.shape[0]
|
||||
head = values.shape[1]
|
||||
channel = values.shape[2]
|
||||
length = values.shape[3]
|
||||
# index init
|
||||
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
||||
# find top k
|
||||
top_k = int(self.factor * math.log(length))
|
||||
weights = torch.topk(corr, top_k, dim=-1)[0]
|
||||
delay = torch.topk(corr, top_k, dim=-1)[1]
|
||||
# update corr
|
||||
tmp_corr = torch.softmax(weights, dim=-1)
|
||||
# aggregation
|
||||
tmp_values = values.repeat(1, 1, 1, 2)
|
||||
delays_agg = torch.zeros_like(values).float()
|
||||
for i in range(top_k):
|
||||
tmp_delay = init_index + delay[..., i].unsqueeze(-1)
|
||||
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
||||
delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
|
||||
return delays_agg
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
if L > S:
|
||||
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
|
||||
values = torch.cat([values, zeros], dim=1)
|
||||
keys = torch.cat([keys, zeros], dim=1)
|
||||
else:
|
||||
values = values[:, :L, :, :]
|
||||
keys = keys[:, :L, :, :]
|
||||
|
||||
# period-based dependencies
|
||||
if self.use_wavelet != 2:
|
||||
if self.use_wavelet == 1:
|
||||
j_list = self.j_list
|
||||
queries = queries.reshape([B, L, -1])
|
||||
keys = keys.reshape([B, L, -1])
|
||||
Ql, Qh_list = self.dwt1d(queries.transpose(1, 2)) # [B, H*D, L]
|
||||
Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
|
||||
qs = [queries.transpose(1, 2)] + Qh_list + [Ql] # [B, H*D, L]
|
||||
ks = [keys.transpose(1, 2)] + Kh_list + [Kl]
|
||||
q_list = []
|
||||
k_list = []
|
||||
for q, k, j in zip(qs, ks, j_list):
|
||||
q_list += [interpolate(q, scale_factor=j, mode='linear')[:, :, -L:]]
|
||||
k_list += [interpolate(k, scale_factor=j, mode='linear')[:, :, -L:]]
|
||||
queries = torch.stack([i.reshape([B, H, E, L]) for i in q_list], dim=3).reshape([B, H, -1, L]).permute(0, 3, 1, 2)
|
||||
keys = torch.stack([i.reshape([B, H, E, L]) for i in k_list], dim=3).reshape([B, H, -1, L]).permute(0, 3, 1, 2)
|
||||
else:
|
||||
pass
|
||||
q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1) # size=[B, H, E, L]
|
||||
k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
|
||||
res = q_fft * torch.conj(k_fft)
|
||||
corr = torch.fft.irfft(res, dim=-1) # size=[B, H, E, L]
|
||||
|
||||
# time delay agg
|
||||
if self.training:
|
||||
V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2) # [B, L, H, E], [B, H, E, L] -> [B, L, H, E]
|
||||
else:
|
||||
V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
||||
else:
|
||||
V_list = []
|
||||
queries = queries.reshape([B, L, -1])
|
||||
keys = keys.reshape([B, L, -1])
|
||||
values = values.reshape([B, L, -1])
|
||||
Ql, Qh_list = self.dwt1d(queries.transpose(1, 2)) # [B, H*D, L]
|
||||
Kl, Kh_list = self.dwt1d(keys.transpose(1, 2))
|
||||
Vl, Vh_list = self.dwt1d(values.transpose(1, 2))
|
||||
qs = Qh_list + [Ql] # [B, H*D, L]
|
||||
ks = Kh_list + [Kl]
|
||||
vs = Vh_list + [Vl]
|
||||
for q, k, v in zip(qs, ks, vs):
|
||||
q = q.reshape([B, H, E, -1])
|
||||
k = k.reshape([B, H, E, -1])
|
||||
v = v.reshape([B, H, E, -1]).permute(0, 3, 1, 2)
|
||||
q_fft = torch.fft.rfft(q.contiguous(), dim=-1)
|
||||
k_fft = torch.fft.rfft(k.contiguous(), dim=-1)
|
||||
res = q_fft * torch.conj(k_fft)
|
||||
corr = torch.fft.irfft(res, dim=-1) # [B, H, E, L]
|
||||
if self.training:
|
||||
V = self.time_delay_agg_training(v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
||||
else:
|
||||
V = self.time_delay_agg_inference(v.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
||||
V_list += [V]
|
||||
Vl = V_list[-1].reshape([B, -1, H*E]).transpose(1, 2)
|
||||
Vh_list = [i.reshape([B, -1, H*E]).transpose(1, 2) for i in V_list[:-1]]
|
||||
V = self.dwt1div((Vl, Vh_list)).reshape([B, H, E, -1]).permute(0, 3, 1, 2)
|
||||
# corr = self.dwt1div((V_list[-1], V_list[:-1]))
|
||||
|
||||
if self.output_attention:
|
||||
return (V.contiguous(), corr.permute(0, 3, 1, 2)) # size = [B, L, H, E]
|
||||
else:
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class AutoCorrelationLayer(nn.Module):
|
||||
def __init__(self, correlation, d_model, n_heads, d_keys=None,
|
||||
d_values=None):
|
||||
super(AutoCorrelationLayer, self).__init__()
|
||||
|
||||
d_keys = d_keys or (d_model // n_heads)
|
||||
d_values = d_values or (d_model // n_heads)
|
||||
|
||||
self.inner_correlation = correlation
|
||||
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.value_projection = nn.Linear(d_model, d_values * n_heads)
|
||||
self.out_projection = nn.Linear(d_values * n_heads, d_model)
|
||||
self.n_heads = n_heads
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, _ = queries.shape
|
||||
_, S, _ = keys.shape
|
||||
H = self.n_heads
|
||||
|
||||
queries = self.query_projection(queries).view(B, L, H, -1)
|
||||
keys = self.key_projection(keys).view(B, S, H, -1)
|
||||
values = self.value_projection(values).view(B, S, H, -1)
|
||||
|
||||
out, attn = self.inner_correlation(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
attn_mask
|
||||
)
|
||||
|
||||
out = out.view(B, L, -1)
|
||||
return self.out_projection(out), attn
|
||||
-219
@@ -1,219 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
from layers.SelfAttention_Family import FullAttention
|
||||
|
||||
|
||||
class my_Layernorm(nn.Module):
|
||||
"""
|
||||
Special designed layernorm for the seasonal part
|
||||
"""
|
||||
def __init__(self, channels):
|
||||
super(my_Layernorm, self).__init__()
|
||||
self.layernorm = nn.LayerNorm(channels)
|
||||
|
||||
def forward(self, x):
|
||||
x_hat = self.layernorm(x)
|
||||
bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
|
||||
return x_hat - bias
|
||||
|
||||
|
||||
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-math.floor((self.kernel_size - 1) // 2), 1)
|
||||
end = x[:, -1:, :].repeat(1, math.floor((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 series_decomp_multi(nn.Module):
|
||||
"""
|
||||
Series decomposition block
|
||||
"""
|
||||
def __init__(self, kernel_size):
|
||||
super(series_decomp_multi, self).__init__()
|
||||
self.moving_avg = [moving_avg(kernel, stride=1) for kernel in kernel_size]
|
||||
self.layer = torch.nn.Linear(1, len(kernel_size))
|
||||
|
||||
def forward(self, x):
|
||||
moving_mean=[]
|
||||
for func in self.moving_avg:
|
||||
moving_avg = func(x)
|
||||
moving_mean.append(moving_avg.unsqueeze(-1))
|
||||
moving_mean=torch.cat(moving_mean,dim=-1)
|
||||
moving_mean = torch.sum(moving_mean*nn.Softmax(-1)(self.layer(x.unsqueeze(-1))),dim=-1)
|
||||
res = x - moving_mean
|
||||
return res, moving_mean
|
||||
|
||||
|
||||
class FourierDecomp(nn.Module):
|
||||
def __init__(self):
|
||||
super(FourierDecomp, self).__init__()
|
||||
pass
|
||||
|
||||
def forward(self, x):
|
||||
x_ft = torch.fft.rfft(x, dim=-1)
|
||||
|
||||
|
||||
class EncoderLayer(nn.Module):
|
||||
"""
|
||||
Autoformer encoder layer with the progressive decomposition architecture
|
||||
"""
|
||||
def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
|
||||
super(EncoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.attention = attention
|
||||
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
|
||||
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
|
||||
|
||||
if isinstance(moving_avg, list):
|
||||
self.decomp1 = series_decomp_multi(moving_avg)
|
||||
self.decomp2 = series_decomp_multi(moving_avg)
|
||||
else:
|
||||
self.decomp1 = series_decomp(moving_avg)
|
||||
self.decomp2 = series_decomp(moving_avg)
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, attn_mask=None):
|
||||
new_x, attn = self.attention(
|
||||
x, x, x,
|
||||
attn_mask=attn_mask
|
||||
)
|
||||
x = x + self.dropout(new_x)
|
||||
x, _ = self.decomp1(x)
|
||||
y = x
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
res, _ = self.decomp2(x + y)
|
||||
return res, attn
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
"""
|
||||
Autoformer encoder
|
||||
"""
|
||||
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):
|
||||
attns = []
|
||||
if self.conv_layers is not None:
|
||||
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
x = conv_layer(x)
|
||||
attns.append(attn)
|
||||
x, attn = self.attn_layers[-1](x)
|
||||
attns.append(attn)
|
||||
else:
|
||||
for attn_layer in self.attn_layers:
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
attns.append(attn)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
return x, attns
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
"""
|
||||
Autoformer decoder layer with the progressive decomposition architecture
|
||||
"""
|
||||
def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
|
||||
moving_avg=25, dropout=0.1, activation="relu"):
|
||||
super(DecoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.self_attention = self_attention
|
||||
self.cross_attention = cross_attention
|
||||
self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
|
||||
self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
|
||||
|
||||
if isinstance(moving_avg, list):
|
||||
self.decomp1 = series_decomp_multi(moving_avg)
|
||||
self.decomp2 = series_decomp_multi(moving_avg)
|
||||
self.decomp3 = series_decomp_multi(moving_avg)
|
||||
else:
|
||||
self.decomp1 = series_decomp(moving_avg)
|
||||
self.decomp2 = series_decomp(moving_avg)
|
||||
self.decomp3 = series_decomp(moving_avg)
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
|
||||
padding_mode='circular', bias=False)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
x = x + self.dropout(self.self_attention(
|
||||
x, x, x,
|
||||
attn_mask=x_mask
|
||||
)[0])
|
||||
|
||||
x, trend1 = self.decomp1(x)
|
||||
x = x + self.dropout(self.cross_attention(
|
||||
x, cross, cross,
|
||||
attn_mask=cross_mask
|
||||
)[0])
|
||||
|
||||
x, trend2 = self.decomp2(x)
|
||||
y = x
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
x, trend3 = self.decomp3(x + y)
|
||||
|
||||
residual_trend = trend1 + trend2 + trend3
|
||||
residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
|
||||
return x, residual_trend
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
"""
|
||||
Autoformer encoder
|
||||
"""
|
||||
def __init__(self, layers, norm_layer=None, projection=None):
|
||||
super(Decoder, self).__init__()
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.norm = norm_layer
|
||||
self.projection = projection
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
|
||||
for layer in self.layers:
|
||||
x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
|
||||
trend = trend + residual_trend
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
if self.projection is not None:
|
||||
x = self.projection(x)
|
||||
return x, trend
|
||||
Vendored
-178
@@ -1,178 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.utils import weight_norm
|
||||
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):
|
||||
x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
|
||||
return self.dropout(x)
|
||||
|
||||
class DataEmbedding_onlypos(nn.Module):
|
||||
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
|
||||
super(DataEmbedding_onlypos, self).__init__()
|
||||
|
||||
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
|
||||
self.position_embedding = PositionalEmbedding(d_model=d_model)
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
|
||||
def forward(self, x, x_mark):
|
||||
x = self.value_embedding(x) + self.position_embedding(x)
|
||||
return self.dropout(x)
|
||||
|
||||
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):
|
||||
# try:
|
||||
x = self.value_embedding(x) + self.temporal_embedding(x_mark)
|
||||
# except:
|
||||
# a = 1
|
||||
return self.dropout(x)
|
||||
|
||||
class DataEmbedding_wo_pos_temp(nn.Module):
|
||||
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
|
||||
super(DataEmbedding_wo_pos_temp, 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):
|
||||
x = self.value_embedding(x)
|
||||
return self.dropout(x)
|
||||
|
||||
class DataEmbedding_wo_temp(nn.Module):
|
||||
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
|
||||
super(DataEmbedding_wo_temp, 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):
|
||||
x = self.value_embedding(x) + self.position_embedding(x)
|
||||
return self.dropout(x)
|
||||
-137
@@ -1,137 +0,0 @@
|
||||
# coding=utf-8
|
||||
# author=maziqing
|
||||
# email=maziqing.mzq@alibaba-inc.com
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def get_frequency_modes(seq_len, modes=64, mode_select_method='random'):
|
||||
"""
|
||||
get modes on frequency domain:
|
||||
'random' means sampling randomly;
|
||||
'else' means sampling the lowest modes;
|
||||
"""
|
||||
modes = min(modes, seq_len//2)
|
||||
if mode_select_method == 'random':
|
||||
index = list(range(0, seq_len // 2))
|
||||
np.random.shuffle(index)
|
||||
index = index[:modes]
|
||||
else:
|
||||
index = list(range(0, modes))
|
||||
index.sort()
|
||||
return index
|
||||
|
||||
|
||||
# ########## fourier layer #############
|
||||
class FourierBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, seq_len, modes=0, mode_select_method='random'):
|
||||
super(FourierBlock, self).__init__()
|
||||
print('fourier enhanced block used!')
|
||||
"""
|
||||
1D Fourier block. It performs representation learning on frequency domain,
|
||||
it does FFT, linear transform, and Inverse FFT.
|
||||
"""
|
||||
# get modes on frequency domain
|
||||
self.index = get_frequency_modes(seq_len, modes=modes, mode_select_method=mode_select_method)
|
||||
print('modes={}, index={}'.format(modes, self.index))
|
||||
|
||||
self.scale = (1 / (in_channels * out_channels))
|
||||
self.weights1 = nn.Parameter(
|
||||
self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index), dtype=torch.cfloat))
|
||||
|
||||
# Complex multiplication
|
||||
def compl_mul1d(self, input, weights):
|
||||
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
|
||||
return torch.einsum("bhi,hio->bho", input, weights)
|
||||
|
||||
def forward(self, q, k, v, mask):
|
||||
# size = [B, L, H, E]
|
||||
B, L, H, E = q.shape
|
||||
x = q.permute(0, 2, 3, 1)
|
||||
# Compute Fourier coefficients
|
||||
x_ft = torch.fft.rfft(x, dim=-1)
|
||||
# Perform Fourier neural operations
|
||||
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=x.device, dtype=torch.cfloat)
|
||||
for wi, i in enumerate(self.index):
|
||||
if i >= x_ft.shape[3] or wi >= out_ft.shape[3]:
|
||||
continue
|
||||
out_ft[:, :, :, wi] = self.compl_mul1d(x_ft[:, :, :, i], self.weights1[:, :, :, wi])
|
||||
# Return to time domain
|
||||
x = torch.fft.irfft(out_ft, n=x.size(-1))
|
||||
return (x, None)
|
||||
|
||||
|
||||
# ########## Fourier Cross Former ####################
|
||||
class FourierCrossAttention(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=64, mode_select_method='random',
|
||||
activation='tanh', policy=0):
|
||||
super(FourierCrossAttention, self).__init__()
|
||||
print(' fourier enhanced cross attention used!')
|
||||
"""
|
||||
1D Fourier Cross Attention layer. It does FFT, linear transform, attention mechanism and Inverse FFT.
|
||||
"""
|
||||
self.activation = activation
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
# get modes for queries and keys (& values) on frequency domain
|
||||
self.index_q = get_frequency_modes(seq_len_q, modes=modes, mode_select_method=mode_select_method)
|
||||
self.index_kv = get_frequency_modes(seq_len_kv, modes=modes, mode_select_method=mode_select_method)
|
||||
|
||||
print('modes_q={}, index_q={}'.format(len(self.index_q), self.index_q))
|
||||
print('modes_kv={}, index_kv={}'.format(len(self.index_kv), self.index_kv))
|
||||
|
||||
self.scale = (1 / (in_channels * out_channels))
|
||||
self.weights1 = nn.Parameter(
|
||||
self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index_q), dtype=torch.cfloat))
|
||||
|
||||
# Complex multiplication
|
||||
def compl_mul1d(self, input, weights):
|
||||
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
|
||||
return torch.einsum("bhi,hio->bho", input, weights)
|
||||
|
||||
def forward(self, q, k, v, mask):
|
||||
# size = [B, L, H, E]
|
||||
B, L, H, E = q.shape
|
||||
xq = q.permute(0, 2, 3, 1) # size = [B, H, E, L]
|
||||
xk = k.permute(0, 2, 3, 1)
|
||||
xv = v.permute(0, 2, 3, 1)
|
||||
|
||||
# Compute Fourier coefficients
|
||||
xq_ft_ = torch.zeros(B, H, E, len(self.index_q), device=xq.device, dtype=torch.cfloat)
|
||||
xq_ft = torch.fft.rfft(xq, dim=-1)
|
||||
for i, j in enumerate(self.index_q):
|
||||
if j >= xq_ft.shape[3]:
|
||||
continue
|
||||
xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
|
||||
xk_ft_ = torch.zeros(B, H, E, len(self.index_kv), device=xq.device, dtype=torch.cfloat)
|
||||
xk_ft = torch.fft.rfft(xk, dim=-1)
|
||||
for i, j in enumerate(self.index_kv):
|
||||
if j >= xk_ft.shape[3]:
|
||||
continue
|
||||
xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
|
||||
|
||||
# perform attention mechanism on frequency domain
|
||||
xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
|
||||
if self.activation == 'tanh':
|
||||
xqk_ft = xqk_ft.tanh()
|
||||
elif self.activation == 'softmax':
|
||||
xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
|
||||
xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
|
||||
else:
|
||||
raise Exception('{} actiation function is not implemented'.format(self.activation))
|
||||
xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
|
||||
xqkvw = torch.einsum("bhex,heox->bhox", xqkv_ft, self.weights1)
|
||||
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=xq.device, dtype=torch.cfloat)
|
||||
for i, j in enumerate(self.index_q):
|
||||
if i >= xqkvw.shape[3] or j >= out_ft.shape[3]:
|
||||
continue
|
||||
out_ft[:, :, :, j] = xqkvw[:, :, :, i]
|
||||
# Return to time domain
|
||||
out = torch.fft.irfft(out_ft / self.in_channels / self.out_channels, n=xq.size(-1))
|
||||
return (out, None)
|
||||
|
||||
|
||||
|
||||
|
||||
-379
@@ -1,379 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from typing import List, Tuple
|
||||
import math
|
||||
from functools import partial
|
||||
from einops import rearrange, reduce, repeat
|
||||
from torch import nn, einsum, diagonal
|
||||
from math import log2, ceil
|
||||
import pdb
|
||||
from utils.masking import LocalMask
|
||||
from layers.utils import get_filter
|
||||
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class MultiWaveletTransform(nn.Module):
|
||||
"""
|
||||
1D multiwavelet block.
|
||||
"""
|
||||
def __init__(self, ich=1, k=8, alpha=16, c=128,
|
||||
nCZ=1, L=0, base='legendre', attention_dropout=0.1):
|
||||
super(MultiWaveletTransform, self).__init__()
|
||||
print('base', base)
|
||||
self.k = k
|
||||
self.c = c
|
||||
self.L = L
|
||||
self.nCZ = nCZ
|
||||
self.Lk0 = nn.Linear(ich, c * k)
|
||||
self.Lk1 = nn.Linear(c * k, ich)
|
||||
self.ich = ich
|
||||
self.MWT_CZ = nn.ModuleList(MWT_CZ1d(k, alpha, L, c, base) for i in range(nCZ))
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
if L > S:
|
||||
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
|
||||
values = torch.cat([values, zeros], dim=1)
|
||||
keys = torch.cat([keys, zeros], dim=1)
|
||||
else:
|
||||
values = values[:, :L, :, :]
|
||||
keys = keys[:, :L, :, :]
|
||||
values = values.view(B, L, -1)
|
||||
|
||||
V = self.Lk0(values).view(B, L, self.c, -1)
|
||||
for i in range(self.nCZ):
|
||||
V = self.MWT_CZ[i](V)
|
||||
if i < self.nCZ - 1:
|
||||
V = F.relu(V)
|
||||
|
||||
V = self.Lk1(V.view(B, L, -1))
|
||||
V = V.view(B, L, -1, D)
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class MultiWaveletCross(nn.Module):
|
||||
"""
|
||||
1D Multiwavelet Cross Attention layer.
|
||||
"""
|
||||
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes, c=64,
|
||||
k=8, ich=512,
|
||||
L=0,
|
||||
base='legendre',
|
||||
mode_select_method='random',
|
||||
initializer=None, activation='tanh',
|
||||
**kwargs):
|
||||
super(MultiWaveletCross, self).__init__()
|
||||
print('base', base)
|
||||
|
||||
self.c = c
|
||||
self.k = k
|
||||
self.L = L
|
||||
H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
|
||||
H0r = H0 @ PHI0
|
||||
G0r = G0 @ PHI0
|
||||
H1r = H1 @ PHI1
|
||||
G1r = G1 @ PHI1
|
||||
|
||||
H0r[np.abs(H0r) < 1e-8] = 0
|
||||
H1r[np.abs(H1r) < 1e-8] = 0
|
||||
G0r[np.abs(G0r) < 1e-8] = 0
|
||||
G1r[np.abs(G1r) < 1e-8] = 0
|
||||
self.max_item = 3
|
||||
|
||||
self.attn1 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
|
||||
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
|
||||
mode_select_method=mode_select_method)
|
||||
self.attn2 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
|
||||
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
|
||||
mode_select_method=mode_select_method)
|
||||
self.attn3 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
|
||||
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
|
||||
mode_select_method=mode_select_method)
|
||||
self.attn4 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, seq_len_q=seq_len_q,
|
||||
seq_len_kv=seq_len_kv, modes=modes, activation=activation,
|
||||
mode_select_method=mode_select_method)
|
||||
self.T0 = nn.Linear(k, k)
|
||||
self.register_buffer('ec_s', torch.Tensor(
|
||||
np.concatenate((H0.T, H1.T), axis=0)))
|
||||
self.register_buffer('ec_d', torch.Tensor(
|
||||
np.concatenate((G0.T, G1.T), axis=0)))
|
||||
|
||||
self.register_buffer('rc_e', torch.Tensor(
|
||||
np.concatenate((H0r, G0r), axis=0)))
|
||||
self.register_buffer('rc_o', torch.Tensor(
|
||||
np.concatenate((H1r, G1r), axis=0)))
|
||||
|
||||
self.Lk = nn.Linear(ich, c * k)
|
||||
self.Lq = nn.Linear(ich, c * k)
|
||||
self.Lv = nn.Linear(ich, c * k)
|
||||
self.out = nn.Linear(c * k, ich)
|
||||
self.modes1 = modes
|
||||
|
||||
def forward(self, q, k, v, mask=None):
|
||||
B, N, H, E = q.shape # (B, N, H, E) torch.Size([3, 768, 8, 2])
|
||||
_, S, _, _ = k.shape # (B, S, H, E) torch.Size([3, 96, 8, 2])
|
||||
|
||||
q = q.view(q.shape[0], q.shape[1], -1)
|
||||
k = k.view(k.shape[0], k.shape[1], -1)
|
||||
v = v.view(v.shape[0], v.shape[1], -1)
|
||||
q = self.Lq(q)
|
||||
q = q.view(q.shape[0], q.shape[1], self.c, self.k)
|
||||
k = self.Lk(k)
|
||||
k = k.view(k.shape[0], k.shape[1], self.c, self.k)
|
||||
v = self.Lv(v)
|
||||
v = v.view(v.shape[0], v.shape[1], self.c, self.k)
|
||||
|
||||
if N > S:
|
||||
zeros = torch.zeros_like(q[:, :(N - S), :]).float()
|
||||
v = torch.cat([v, zeros], dim=1)
|
||||
k = torch.cat([k, zeros], dim=1)
|
||||
else:
|
||||
v = v[:, :N, :, :]
|
||||
k = k[:, :N, :, :]
|
||||
|
||||
ns = math.floor(np.log2(N))
|
||||
nl = pow(2, math.ceil(np.log2(N)))
|
||||
extra_q = q[:, 0:nl - N, :, :]
|
||||
extra_k = k[:, 0:nl - N, :, :]
|
||||
extra_v = v[:, 0:nl - N, :, :]
|
||||
q = torch.cat([q, extra_q], 1)
|
||||
k = torch.cat([k, extra_k], 1)
|
||||
v = torch.cat([v, extra_v], 1)
|
||||
|
||||
Ud_q = torch.jit.annotate(List[Tuple[Tensor]], [])
|
||||
Ud_k = torch.jit.annotate(List[Tuple[Tensor]], [])
|
||||
Ud_v = torch.jit.annotate(List[Tuple[Tensor]], [])
|
||||
|
||||
Us_q = torch.jit.annotate(List[Tensor], [])
|
||||
Us_k = torch.jit.annotate(List[Tensor], [])
|
||||
Us_v = torch.jit.annotate(List[Tensor], [])
|
||||
|
||||
Ud = torch.jit.annotate(List[Tensor], [])
|
||||
Us = torch.jit.annotate(List[Tensor], [])
|
||||
|
||||
# decompose
|
||||
for i in range(ns - self.L):
|
||||
# print('q shape',q.shape)
|
||||
d, q = self.wavelet_transform(q)
|
||||
Ud_q += [tuple([d, q])]
|
||||
Us_q += [d]
|
||||
for i in range(ns - self.L):
|
||||
d, k = self.wavelet_transform(k)
|
||||
Ud_k += [tuple([d, k])]
|
||||
Us_k += [d]
|
||||
for i in range(ns - self.L):
|
||||
d, v = self.wavelet_transform(v)
|
||||
Ud_v += [tuple([d, v])]
|
||||
Us_v += [d]
|
||||
for i in range(ns - self.L):
|
||||
dk, sk = Ud_k[i], Us_k[i]
|
||||
dq, sq = Ud_q[i], Us_q[i]
|
||||
dv, sv = Ud_v[i], Us_v[i]
|
||||
Ud += [self.attn1(dq[0], dk[0], dv[0], mask)[0] + self.attn2(dq[1], dk[1], dv[1], mask)[0]]
|
||||
Us += [self.attn3(sq, sk, sv, mask)[0]]
|
||||
v = self.attn4(q, k, v, mask)[0]
|
||||
|
||||
# reconstruct
|
||||
for i in range(ns - 1 - self.L, -1, -1):
|
||||
v = v + Us[i]
|
||||
v = torch.cat((v, Ud[i]), -1)
|
||||
v = self.evenOdd(v)
|
||||
v = self.out(v[:, :N, :, :].contiguous().view(B, N, -1))
|
||||
return (v.contiguous(), None)
|
||||
|
||||
def wavelet_transform(self, x):
|
||||
xa = torch.cat([x[:, ::2, :, :],
|
||||
x[:, 1::2, :, :],
|
||||
], -1)
|
||||
d = torch.matmul(xa, self.ec_d)
|
||||
s = torch.matmul(xa, self.ec_s)
|
||||
return d, s
|
||||
|
||||
def evenOdd(self, x):
|
||||
B, N, c, ich = x.shape # (B, N, c, k)
|
||||
assert ich == 2 * self.k
|
||||
x_e = torch.matmul(x, self.rc_e)
|
||||
x_o = torch.matmul(x, self.rc_o)
|
||||
|
||||
x = torch.zeros(B, N * 2, c, self.k,
|
||||
device=x.device)
|
||||
x[..., ::2, :, :] = x_e
|
||||
x[..., 1::2, :, :] = x_o
|
||||
return x
|
||||
|
||||
|
||||
class FourierCrossAttentionW(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=16, activation='tanh',
|
||||
mode_select_method='random'):
|
||||
super(FourierCrossAttentionW, self).__init__()
|
||||
print('corss fourier correlation used!')
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.modes1 = modes
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, q, k, v, mask):
|
||||
B, L, E, H = q.shape
|
||||
|
||||
xq = q.permute(0, 3, 2, 1) # size = [B, H, E, L] torch.Size([3, 8, 64, 512])
|
||||
xk = k.permute(0, 3, 2, 1)
|
||||
xv = v.permute(0, 3, 2, 1)
|
||||
self.index_q = list(range(0, min(int(L // 2), self.modes1)))
|
||||
self.index_k_v = list(range(0, min(int(xv.shape[3] // 2), self.modes1)))
|
||||
|
||||
# Compute Fourier coefficients
|
||||
xq_ft_ = torch.zeros(B, H, E, len(self.index_q), device=xq.device, dtype=torch.cfloat)
|
||||
xq_ft = torch.fft.rfft(xq, dim=-1)
|
||||
for i, j in enumerate(self.index_q):
|
||||
xq_ft_[:, :, :, i] = xq_ft[:, :, :, j]
|
||||
|
||||
xk_ft_ = torch.zeros(B, H, E, len(self.index_k_v), device=xq.device, dtype=torch.cfloat)
|
||||
xk_ft = torch.fft.rfft(xk, dim=-1)
|
||||
for i, j in enumerate(self.index_k_v):
|
||||
xk_ft_[:, :, :, i] = xk_ft[:, :, :, j]
|
||||
xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_))
|
||||
if self.activation == 'tanh':
|
||||
xqk_ft = xqk_ft.tanh()
|
||||
elif self.activation == 'softmax':
|
||||
xqk_ft = torch.softmax(abs(xqk_ft), dim=-1)
|
||||
xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft))
|
||||
else:
|
||||
raise Exception('{} actiation function is not implemented'.format(self.activation))
|
||||
xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_)
|
||||
|
||||
xqkvw = xqkv_ft
|
||||
out_ft = torch.zeros(B, H, E, L // 2 + 1, device=xq.device, dtype=torch.cfloat)
|
||||
for i, j in enumerate(self.index_q):
|
||||
out_ft[:, :, :, j] = xqkvw[:, :, :, i]
|
||||
|
||||
out = torch.fft.irfft(out_ft / self.in_channels / self.out_channels, n=xq.size(-1)).permute(0, 3, 2, 1)
|
||||
# size = [B, L, H, E]
|
||||
return (out, None)
|
||||
|
||||
|
||||
class sparseKernelFT1d(nn.Module):
|
||||
def __init__(self,
|
||||
k, alpha, c=1,
|
||||
nl=1,
|
||||
initializer=None,
|
||||
**kwargs):
|
||||
super(sparseKernelFT1d, self).__init__()
|
||||
|
||||
self.modes1 = alpha
|
||||
self.scale = (1 / (c * k * c * k))
|
||||
self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.cfloat))
|
||||
self.weights1.requires_grad = True
|
||||
self.k = k
|
||||
|
||||
def compl_mul1d(self, x, weights):
|
||||
# (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
|
||||
return torch.einsum("bix,iox->box", x, weights)
|
||||
|
||||
def forward(self, x):
|
||||
B, N, c, k = x.shape # (B, N, c, k)
|
||||
|
||||
x = x.view(B, N, -1)
|
||||
x = x.permute(0, 2, 1)
|
||||
x_fft = torch.fft.rfft(x)
|
||||
# Multiply relevant Fourier modes
|
||||
l = min(self.modes1, N // 2 + 1)
|
||||
# l = N//2+1
|
||||
out_ft = torch.zeros(B, c * k, N // 2 + 1, device=x.device, dtype=torch.cfloat)
|
||||
out_ft[:, :, :l] = self.compl_mul1d(x_fft[:, :, :l], self.weights1[:, :, :l])
|
||||
x = torch.fft.irfft(out_ft, n=N)
|
||||
x = x.permute(0, 2, 1).view(B, N, c, k)
|
||||
return x
|
||||
|
||||
|
||||
# ##
|
||||
class MWT_CZ1d(nn.Module):
|
||||
def __init__(self,
|
||||
k=3, alpha=64,
|
||||
L=0, c=1,
|
||||
base='legendre',
|
||||
initializer=None,
|
||||
**kwargs):
|
||||
super(MWT_CZ1d, self).__init__()
|
||||
|
||||
self.k = k
|
||||
self.L = L
|
||||
H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k)
|
||||
H0r = H0 @ PHI0
|
||||
G0r = G0 @ PHI0
|
||||
H1r = H1 @ PHI1
|
||||
G1r = G1 @ PHI1
|
||||
|
||||
H0r[np.abs(H0r) < 1e-8] = 0
|
||||
H1r[np.abs(H1r) < 1e-8] = 0
|
||||
G0r[np.abs(G0r) < 1e-8] = 0
|
||||
G1r[np.abs(G1r) < 1e-8] = 0
|
||||
self.max_item = 3
|
||||
|
||||
self.A = sparseKernelFT1d(k, alpha, c)
|
||||
self.B = sparseKernelFT1d(k, alpha, c)
|
||||
self.C = sparseKernelFT1d(k, alpha, c)
|
||||
|
||||
self.T0 = nn.Linear(k, k)
|
||||
|
||||
self.register_buffer('ec_s', torch.Tensor(
|
||||
np.concatenate((H0.T, H1.T), axis=0)))
|
||||
self.register_buffer('ec_d', torch.Tensor(
|
||||
np.concatenate((G0.T, G1.T), axis=0)))
|
||||
|
||||
self.register_buffer('rc_e', torch.Tensor(
|
||||
np.concatenate((H0r, G0r), axis=0)))
|
||||
self.register_buffer('rc_o', torch.Tensor(
|
||||
np.concatenate((H1r, G1r), axis=0)))
|
||||
|
||||
def forward(self, x):
|
||||
B, N, c, k = x.shape # (B, N, k)
|
||||
ns = math.floor(np.log2(N))
|
||||
nl = pow(2, math.ceil(np.log2(N)))
|
||||
extra_x = x[:, 0:nl - N, :, :]
|
||||
x = torch.cat([x, extra_x], 1)
|
||||
Ud = torch.jit.annotate(List[Tensor], [])
|
||||
Us = torch.jit.annotate(List[Tensor], [])
|
||||
# decompose
|
||||
for i in range(ns - self.L):
|
||||
# print('x shape',x.shape)
|
||||
d, x = self.wavelet_transform(x)
|
||||
Ud += [self.A(d) + self.B(x)]
|
||||
Us += [self.C(d)]
|
||||
x = self.T0(x) # coarsest scale transform
|
||||
|
||||
# reconstruct
|
||||
for i in range(ns - 1 - self.L, -1, -1):
|
||||
x = x + Us[i]
|
||||
x = torch.cat((x, Ud[i]), -1)
|
||||
x = self.evenOdd(x)
|
||||
x = x[:, :N, :, :]
|
||||
|
||||
return x
|
||||
|
||||
def wavelet_transform(self, x):
|
||||
xa = torch.cat([x[:, ::2, :, :],
|
||||
x[:, 1::2, :, :],
|
||||
], -1)
|
||||
d = torch.matmul(xa, self.ec_d)
|
||||
s = torch.matmul(xa, self.ec_s)
|
||||
return d, s
|
||||
|
||||
def evenOdd(self, x):
|
||||
|
||||
B, N, c, ich = x.shape # (B, N, c, k)
|
||||
assert ich == 2 * self.k
|
||||
x_e = torch.matmul(x, self.rc_e)
|
||||
x_o = torch.matmul(x, self.rc_o)
|
||||
|
||||
x = torch.zeros(B, N * 2, c, self.k,
|
||||
device=x.device)
|
||||
x[..., ::2, :, :] = x_e
|
||||
x[..., 1::2, :, :] = x_o
|
||||
return x
|
||||
-198
@@ -1,198 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
import numpy as np
|
||||
import math
|
||||
from math import sqrt
|
||||
from utils.masking import TriangularCausalMask, ProbMask
|
||||
import os
|
||||
|
||||
|
||||
class FullAttention(nn.Module):
|
||||
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(FullAttention, self).__init__()
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
scale = self.scale or 1. / sqrt(E)
|
||||
# print('queries shape',queries.shape)
|
||||
# print('keys shape',keys.shape)
|
||||
# print('values shape',values.shape)
|
||||
|
||||
scores = torch.einsum("blhe,bshe->bhls", queries, keys)
|
||||
|
||||
if self.mask_flag:
|
||||
if attn_mask is None:
|
||||
attn_mask = TriangularCausalMask(B, L, device=queries.device)
|
||||
|
||||
scores.masked_fill_(attn_mask.mask, -np.inf)
|
||||
|
||||
A = self.dropout(torch.softmax(scale * scores, dim=-1))
|
||||
V = torch.einsum("bhls,bshd->blhd", A, values)
|
||||
#print('output shape',V.shape)
|
||||
|
||||
if self.output_attention:
|
||||
return (V.contiguous(), A)
|
||||
else:
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class SparseAttention(nn.Module):
|
||||
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(FullAttention, self).__init__()
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask=None):
|
||||
B, L, H, E = queries.shape
|
||||
_, S, _, D = values.shape
|
||||
scale = self.scale or 1. / sqrt(E)
|
||||
scores = torch.einsum("blhe,bshe->bhls", queries, keys)
|
||||
if attn_mask is None:
|
||||
attn_mask = LocalMask(B, L, device=queries.device)
|
||||
|
||||
scores.masked_fill_(attn_mask.mask, -np.inf)
|
||||
|
||||
A = self.dropout(torch.softmax(scale * scores, dim=-1))
|
||||
V = torch.einsum("bhls,bshd->blhd", A, values)
|
||||
#print('output shape',V.shape)
|
||||
|
||||
if self.output_attention:
|
||||
return (V.contiguous(), A)
|
||||
else:
|
||||
return (V.contiguous(), None)
|
||||
|
||||
|
||||
class ProbAttention(nn.Module):
|
||||
def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
|
||||
super(ProbAttention, self).__init__()
|
||||
self.factor = factor
|
||||
self.scale = scale
|
||||
self.mask_flag = mask_flag
|
||||
self.output_attention = output_attention
|
||||
self.dropout = nn.Dropout(attention_dropout)
|
||||
|
||||
def _prob_QK(self, Q, K, sample_k, n_top): # n_top: c*ln(L_q)
|
||||
# Q [B, H, L, D]
|
||||
B, H, L_K, E = K.shape
|
||||
_, _, L_Q, _ = Q.shape
|
||||
|
||||
# calculate the sampled Q_K
|
||||
K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
|
||||
index_sample = torch.randint(L_K, (L_Q, sample_k)) # real U = U_part(factor*ln(L_k))*L_q
|
||||
K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]
|
||||
Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
|
||||
|
||||
# find the Top_k query with sparisty measurement
|
||||
M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
|
||||
M_top = M.topk(n_top, sorted=False)[1]
|
||||
|
||||
# use the reduced Q to calculate Q_K
|
||||
Q_reduce = Q[torch.arange(B)[:, None, None],
|
||||
torch.arange(H)[None, :, None],
|
||||
M_top, :] # factor*ln(L_q)
|
||||
Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1)) # factor*ln(L_q)*L_k
|
||||
|
||||
return Q_K, M_top
|
||||
|
||||
def _get_initial_context(self, V, L_Q):
|
||||
B, H, L_V, D = V.shape
|
||||
if not self.mask_flag:
|
||||
# V_sum = V.sum(dim=-2)
|
||||
V_sum = V.mean(dim=-2)
|
||||
contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
|
||||
else: # use mask
|
||||
assert (L_Q == L_V) # requires that L_Q == L_V, i.e. for self-attention only
|
||||
contex = V.cumsum(dim=-2)
|
||||
return contex
|
||||
|
||||
def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
|
||||
B, H, L_V, D = V.shape
|
||||
|
||||
if self.mask_flag:
|
||||
attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
|
||||
scores.masked_fill_(attn_mask.mask, -np.inf)
|
||||
|
||||
attn = torch.softmax(scores, dim=-1) # nn.Softmax(dim=-1)(scores)
|
||||
|
||||
context_in[torch.arange(B)[:, None, None],
|
||||
torch.arange(H)[None, :, None],
|
||||
index, :] = torch.matmul(attn, V).type_as(context_in)
|
||||
if self.output_attention:
|
||||
attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)
|
||||
attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn
|
||||
return (context_in, attns)
|
||||
else:
|
||||
return (context_in, None)
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L_Q, H, D = queries.shape
|
||||
_, L_K, _, _ = keys.shape
|
||||
|
||||
queries = queries.transpose(2, 1)
|
||||
keys = keys.transpose(2, 1)
|
||||
values = values.transpose(2, 1)
|
||||
|
||||
U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item() # c*ln(L_k)
|
||||
u = self.factor * np.ceil(np.log(L_Q)).astype('int').item() # c*ln(L_q)
|
||||
|
||||
U_part = U_part if U_part < L_K else L_K
|
||||
u = u if u < L_Q else L_Q
|
||||
|
||||
scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
|
||||
|
||||
# add scale factor
|
||||
scale = self.scale or 1. / sqrt(D)
|
||||
if scale is not None:
|
||||
scores_top = scores_top * scale
|
||||
# get the context
|
||||
context = self._get_initial_context(values, L_Q)
|
||||
# update the context with selected top_k queries
|
||||
context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
|
||||
|
||||
return context.contiguous(), attn
|
||||
|
||||
|
||||
class AttentionLayer(nn.Module):
|
||||
def __init__(self, attention, d_model, n_heads, d_keys=None,
|
||||
d_values=None):
|
||||
super(AttentionLayer, self).__init__()
|
||||
|
||||
d_keys = d_keys or (d_model // n_heads)
|
||||
d_values = d_values or (d_model // n_heads)
|
||||
|
||||
self.inner_attention = attention
|
||||
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
|
||||
self.value_projection = nn.Linear(d_model, d_values * n_heads)
|
||||
self.out_projection = nn.Linear(d_values * n_heads, d_model)
|
||||
self.n_heads = n_heads
|
||||
|
||||
def forward(self, queries, keys, values, attn_mask):
|
||||
B, L, _ = queries.shape
|
||||
_, S, _ = keys.shape
|
||||
H = self.n_heads
|
||||
|
||||
queries = self.query_projection(queries).view(B, L, H, -1)
|
||||
keys = self.key_projection(keys).view(B, S, H, -1)
|
||||
values = self.value_projection(values).view(B, S, H, -1)
|
||||
|
||||
out, attn = self.inner_attention(
|
||||
queries,
|
||||
keys,
|
||||
values,
|
||||
attn_mask
|
||||
)
|
||||
out = out.view(B, L, -1)
|
||||
|
||||
return self.out_projection(out), attn
|
||||
-131
@@ -1,131 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(self, c_in):
|
||||
super(ConvLayer, self).__init__()
|
||||
self.downConv = nn.Conv1d(in_channels=c_in,
|
||||
out_channels=c_in,
|
||||
kernel_size=3,
|
||||
padding=2,
|
||||
padding_mode='circular')
|
||||
self.norm = nn.BatchNorm1d(c_in)
|
||||
self.activation = nn.ELU()
|
||||
self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.downConv(x.permute(0, 2, 1))
|
||||
x = self.norm(x)
|
||||
x = self.activation(x)
|
||||
x = self.maxPool(x)
|
||||
x = x.transpose(1, 2)
|
||||
return x
|
||||
|
||||
|
||||
class EncoderLayer(nn.Module):
|
||||
def __init__(self, attention, 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.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
|
||||
|
||||
def forward(self, x, attn_mask=None):
|
||||
new_x, attn = self.attention(
|
||||
x, x, x,
|
||||
attn_mask=attn_mask
|
||||
)
|
||||
x = x + self.dropout(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):
|
||||
# x [B, L, D]
|
||||
attns = []
|
||||
if self.conv_layers is not None:
|
||||
for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
x = conv_layer(x)
|
||||
attns.append(attn)
|
||||
x, attn = self.attn_layers[-1](x)
|
||||
attns.append(attn)
|
||||
else:
|
||||
for attn_layer in self.attn_layers:
|
||||
x, attn = attn_layer(x, attn_mask=attn_mask)
|
||||
attns.append(attn)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
return x, attns
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
|
||||
dropout=0.1, activation="relu"):
|
||||
super(DecoderLayer, self).__init__()
|
||||
d_ff = d_ff or 4 * d_model
|
||||
self.self_attention = self_attention
|
||||
self.cross_attention = cross_attention
|
||||
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.norm3 = nn.LayerNorm(d_model)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.activation = F.relu if activation == "relu" else F.gelu
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
x = x + self.dropout(self.self_attention(
|
||||
x, x, x,
|
||||
attn_mask=x_mask
|
||||
)[0])
|
||||
x = self.norm1(x)
|
||||
|
||||
x = x + self.dropout(self.cross_attention(
|
||||
x, cross, cross,
|
||||
attn_mask=cross_mask
|
||||
)[0])
|
||||
|
||||
y = x = self.norm2(x)
|
||||
y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
|
||||
y = self.dropout(self.conv2(y).transpose(-1, 1))
|
||||
|
||||
return self.norm3(x + y)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, layers, norm_layer=None, projection=None):
|
||||
super(Decoder, self).__init__()
|
||||
self.layers = nn.ModuleList(layers)
|
||||
self.norm = norm_layer
|
||||
self.projection = projection
|
||||
|
||||
def forward(self, x, cross, x_mask=None, cross_mask=None):
|
||||
for layer in self.layers:
|
||||
x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
|
||||
|
||||
if self.norm is not None:
|
||||
x = self.norm(x)
|
||||
|
||||
if self.projection is not None:
|
||||
x = self.projection(x)
|
||||
return x
|
||||
Vendored
-389
@@ -1,389 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import numpy as np
|
||||
from functools import partial
|
||||
|
||||
from scipy.special import eval_legendre
|
||||
from sympy import Poly, legendre, Symbol, chebyshevt
|
||||
|
||||
def legendreDer(k, x):
|
||||
def _legendre(k, x):
|
||||
return (2*k+1) * eval_legendre(k, x)
|
||||
out = 0
|
||||
for i in np.arange(k-1,-1,-2):
|
||||
out += _legendre(i, x)
|
||||
return out
|
||||
|
||||
def phi_(phi_c, x, lb = 0, ub = 1):
|
||||
mask = np.logical_or(x<lb, x>ub) * 1.0
|
||||
return np.polynomial.polynomial.Polynomial(phi_c)(x) * (1-mask)
|
||||
|
||||
def get_phi_psi(k, base):
|
||||
|
||||
x = Symbol('x')
|
||||
phi_coeff = np.zeros((k,k))
|
||||
phi_2x_coeff = np.zeros((k,k))
|
||||
if base == 'legendre':
|
||||
for ki in range(k):
|
||||
coeff_ = Poly(legendre(ki, 2*x-1), x).all_coeffs()
|
||||
phi_coeff[ki,:ki+1] = np.flip(np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64))
|
||||
coeff_ = Poly(legendre(ki, 4*x-1), x).all_coeffs()
|
||||
phi_2x_coeff[ki,:ki+1] = np.flip(np.sqrt(2) * np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64))
|
||||
|
||||
psi1_coeff = np.zeros((k, k))
|
||||
psi2_coeff = np.zeros((k, k))
|
||||
for ki in range(k):
|
||||
psi1_coeff[ki,:] = phi_2x_coeff[ki,:]
|
||||
for i in range(k):
|
||||
a = phi_2x_coeff[ki,:ki+1]
|
||||
b = phi_coeff[i, :i+1]
|
||||
prod_ = np.convolve(a, b)
|
||||
prod_[np.abs(prod_)<1e-8] = 0
|
||||
proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
|
||||
psi1_coeff[ki,:] -= proj_ * phi_coeff[i,:]
|
||||
psi2_coeff[ki,:] -= proj_ * phi_coeff[i,:]
|
||||
for j in range(ki):
|
||||
a = phi_2x_coeff[ki,:ki+1]
|
||||
b = psi1_coeff[j, :]
|
||||
prod_ = np.convolve(a, b)
|
||||
prod_[np.abs(prod_)<1e-8] = 0
|
||||
proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
|
||||
psi1_coeff[ki,:] -= proj_ * psi1_coeff[j,:]
|
||||
psi2_coeff[ki,:] -= proj_ * psi2_coeff[j,:]
|
||||
|
||||
a = psi1_coeff[ki,:]
|
||||
prod_ = np.convolve(a, a)
|
||||
prod_[np.abs(prod_)<1e-8] = 0
|
||||
norm1 = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum()
|
||||
|
||||
a = psi2_coeff[ki,:]
|
||||
prod_ = np.convolve(a, a)
|
||||
prod_[np.abs(prod_)<1e-8] = 0
|
||||
norm2 = (prod_ * 1/(np.arange(len(prod_))+1) * (1-np.power(0.5, 1+np.arange(len(prod_))))).sum()
|
||||
norm_ = np.sqrt(norm1 + norm2)
|
||||
psi1_coeff[ki,:] /= norm_
|
||||
psi2_coeff[ki,:] /= norm_
|
||||
psi1_coeff[np.abs(psi1_coeff)<1e-8] = 0
|
||||
psi2_coeff[np.abs(psi2_coeff)<1e-8] = 0
|
||||
|
||||
phi = [np.poly1d(np.flip(phi_coeff[i,:])) for i in range(k)]
|
||||
psi1 = [np.poly1d(np.flip(psi1_coeff[i,:])) for i in range(k)]
|
||||
psi2 = [np.poly1d(np.flip(psi2_coeff[i,:])) for i in range(k)]
|
||||
|
||||
elif base == 'chebyshev':
|
||||
for ki in range(k):
|
||||
if ki == 0:
|
||||
phi_coeff[ki,:ki+1] = np.sqrt(2/np.pi)
|
||||
phi_2x_coeff[ki,:ki+1] = np.sqrt(2/np.pi) * np.sqrt(2)
|
||||
else:
|
||||
coeff_ = Poly(chebyshevt(ki, 2*x-1), x).all_coeffs()
|
||||
phi_coeff[ki,:ki+1] = np.flip(2/np.sqrt(np.pi) * np.array(coeff_).astype(np.float64))
|
||||
coeff_ = Poly(chebyshevt(ki, 4*x-1), x).all_coeffs()
|
||||
phi_2x_coeff[ki,:ki+1] = np.flip(np.sqrt(2) * 2 / np.sqrt(np.pi) * np.array(coeff_).astype(np.float64))
|
||||
|
||||
phi = [partial(phi_, phi_coeff[i,:]) for i in range(k)]
|
||||
|
||||
x = Symbol('x')
|
||||
kUse = 2*k
|
||||
roots = Poly(chebyshevt(kUse, 2*x-1)).all_roots()
|
||||
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
|
||||
# x_m[x_m==0.5] = 0.5 + 1e-8 # add small noise to avoid the case of 0.5 belonging to both phi(2x) and phi(2x-1)
|
||||
# not needed for our purpose here, we use even k always to avoid
|
||||
wm = np.pi / kUse / 2
|
||||
|
||||
psi1_coeff = np.zeros((k, k))
|
||||
psi2_coeff = np.zeros((k, k))
|
||||
|
||||
psi1 = [[] for _ in range(k)]
|
||||
psi2 = [[] for _ in range(k)]
|
||||
|
||||
for ki in range(k):
|
||||
psi1_coeff[ki,:] = phi_2x_coeff[ki,:]
|
||||
for i in range(k):
|
||||
proj_ = (wm * phi[i](x_m) * np.sqrt(2)* phi[ki](2*x_m)).sum()
|
||||
psi1_coeff[ki,:] -= proj_ * phi_coeff[i,:]
|
||||
psi2_coeff[ki,:] -= proj_ * phi_coeff[i,:]
|
||||
|
||||
for j in range(ki):
|
||||
proj_ = (wm * psi1[j](x_m) * np.sqrt(2) * phi[ki](2*x_m)).sum()
|
||||
psi1_coeff[ki,:] -= proj_ * psi1_coeff[j,:]
|
||||
psi2_coeff[ki,:] -= proj_ * psi2_coeff[j,:]
|
||||
|
||||
psi1[ki] = partial(phi_, psi1_coeff[ki,:], lb = 0, ub = 0.5)
|
||||
psi2[ki] = partial(phi_, psi2_coeff[ki,:], lb = 0.5, ub = 1)
|
||||
|
||||
norm1 = (wm * psi1[ki](x_m) * psi1[ki](x_m)).sum()
|
||||
norm2 = (wm * psi2[ki](x_m) * psi2[ki](x_m)).sum()
|
||||
|
||||
norm_ = np.sqrt(norm1 + norm2)
|
||||
psi1_coeff[ki,:] /= norm_
|
||||
psi2_coeff[ki,:] /= norm_
|
||||
psi1_coeff[np.abs(psi1_coeff)<1e-8] = 0
|
||||
psi2_coeff[np.abs(psi2_coeff)<1e-8] = 0
|
||||
|
||||
psi1[ki] = partial(phi_, psi1_coeff[ki,:], lb = 0, ub = 0.5+1e-16)
|
||||
psi2[ki] = partial(phi_, psi2_coeff[ki,:], lb = 0.5+1e-16, ub = 1)
|
||||
|
||||
return phi, psi1, psi2
|
||||
|
||||
|
||||
def get_filter(base, k):
|
||||
|
||||
def psi(psi1, psi2, i, inp):
|
||||
mask = (inp<=0.5) * 1.0
|
||||
return psi1[i](inp) * mask + psi2[i](inp) * (1-mask)
|
||||
|
||||
if base not in ['legendre', 'chebyshev']:
|
||||
raise Exception('Base not supported')
|
||||
|
||||
x = Symbol('x')
|
||||
H0 = np.zeros((k,k))
|
||||
H1 = np.zeros((k,k))
|
||||
G0 = np.zeros((k,k))
|
||||
G1 = np.zeros((k,k))
|
||||
PHI0 = np.zeros((k,k))
|
||||
PHI1 = np.zeros((k,k))
|
||||
phi, psi1, psi2 = get_phi_psi(k, base)
|
||||
if base == 'legendre':
|
||||
roots = Poly(legendre(k, 2*x-1)).all_roots()
|
||||
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
|
||||
wm = 1/k/legendreDer(k,2*x_m-1)/eval_legendre(k-1,2*x_m-1)
|
||||
|
||||
for ki in range(k):
|
||||
for kpi in range(k):
|
||||
H0[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki](x_m/2) * phi[kpi](x_m)).sum()
|
||||
G0[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, x_m/2) * phi[kpi](x_m)).sum()
|
||||
H1[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki]((x_m+1)/2) * phi[kpi](x_m)).sum()
|
||||
G1[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, (x_m+1)/2) * phi[kpi](x_m)).sum()
|
||||
|
||||
PHI0 = np.eye(k)
|
||||
PHI1 = np.eye(k)
|
||||
|
||||
elif base == 'chebyshev':
|
||||
x = Symbol('x')
|
||||
kUse = 2*k
|
||||
roots = Poly(chebyshevt(kUse, 2*x-1)).all_roots()
|
||||
x_m = np.array([rt.evalf(20) for rt in roots]).astype(np.float64)
|
||||
# x_m[x_m==0.5] = 0.5 + 1e-8 # add small noise to avoid the case of 0.5 belonging to both phi(2x) and phi(2x-1)
|
||||
# not needed for our purpose here, we use even k always to avoid
|
||||
wm = np.pi / kUse / 2
|
||||
|
||||
for ki in range(k):
|
||||
for kpi in range(k):
|
||||
H0[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki](x_m/2) * phi[kpi](x_m)).sum()
|
||||
G0[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, x_m/2) * phi[kpi](x_m)).sum()
|
||||
H1[ki, kpi] = 1/np.sqrt(2) * (wm * phi[ki]((x_m+1)/2) * phi[kpi](x_m)).sum()
|
||||
G1[ki, kpi] = 1/np.sqrt(2) * (wm * psi(psi1, psi2, ki, (x_m+1)/2) * phi[kpi](x_m)).sum()
|
||||
|
||||
PHI0[ki, kpi] = (wm * phi[ki](2*x_m) * phi[kpi](2*x_m)).sum() * 2
|
||||
PHI1[ki, kpi] = (wm * phi[ki](2*x_m-1) * phi[kpi](2*x_m-1)).sum() * 2
|
||||
|
||||
PHI0[np.abs(PHI0)<1e-8] = 0
|
||||
PHI1[np.abs(PHI1)<1e-8] = 0
|
||||
|
||||
H0[np.abs(H0)<1e-8] = 0
|
||||
H1[np.abs(H1)<1e-8] = 0
|
||||
G0[np.abs(G0)<1e-8] = 0
|
||||
G1[np.abs(G1)<1e-8] = 0
|
||||
|
||||
return H0, H1, G0, G1, PHI0, PHI1
|
||||
|
||||
|
||||
def train(model, train_loader, optimizer, epoch, device, verbose = 0,
|
||||
lossFn = None, lr_schedule=None,
|
||||
post_proc = lambda args: args):
|
||||
|
||||
if lossFn is None:
|
||||
lossFn = nn.MSELoss()
|
||||
|
||||
model.train()
|
||||
|
||||
total_loss = 0.
|
||||
|
||||
for batch_idx, (data, target) in enumerate(train_loader):
|
||||
|
||||
bs = len(data)
|
||||
data, target = data.to(device), target.to(device)
|
||||
optimizer.zero_grad()
|
||||
|
||||
output = model(data)
|
||||
|
||||
target = post_proc(target)
|
||||
output = post_proc(output)
|
||||
loss = lossFn(output.view(bs, -1), target.view(bs, -1))
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.sum().item()
|
||||
if lr_schedule is not None: lr_schedule.step()
|
||||
|
||||
if verbose>0:
|
||||
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
|
||||
epoch, batch_idx * len(data), len(train_loader.dataset),
|
||||
100. * batch_idx / len(train_loader), loss.item()))
|
||||
|
||||
return total_loss/len(train_loader.dataset)
|
||||
|
||||
|
||||
def test(model, test_loader, device, verbose=0, lossFn=None,
|
||||
post_proc = lambda args: args):
|
||||
|
||||
model.eval()
|
||||
if lossFn is None:
|
||||
lossFn = nn.MSELoss()
|
||||
|
||||
|
||||
total_loss = 0.
|
||||
predictions = []
|
||||
|
||||
with torch.no_grad():
|
||||
for data, target in test_loader:
|
||||
bs = len(data)
|
||||
|
||||
data, target = data.to(device), target.to(device)
|
||||
output = model(data)
|
||||
output = post_proc(output)
|
||||
|
||||
loss = lossFn(output.view(bs, -1), target.view(bs, -1))
|
||||
total_loss += loss.sum().item()
|
||||
|
||||
return total_loss/len(test_loader.dataset)
|
||||
|
||||
|
||||
# Till EoF
|
||||
# taken from FNO paper:
|
||||
# https://github.com/zongyi-li/fourier_neural_operator
|
||||
|
||||
# normalization, pointwise gaussian
|
||||
class UnitGaussianNormalizer(object):
|
||||
def __init__(self, x, eps=0.00001):
|
||||
super(UnitGaussianNormalizer, self).__init__()
|
||||
|
||||
# x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T
|
||||
self.mean = torch.mean(x, 0)
|
||||
self.std = torch.std(x, 0)
|
||||
self.eps = eps
|
||||
|
||||
def encode(self, x):
|
||||
x = (x - self.mean) / (self.std + self.eps)
|
||||
return x
|
||||
|
||||
def decode(self, x, sample_idx=None):
|
||||
if sample_idx is None:
|
||||
std = self.std + self.eps # n
|
||||
mean = self.mean
|
||||
else:
|
||||
if len(self.mean.shape) == len(sample_idx[0].shape):
|
||||
std = self.std[sample_idx] + self.eps # batch*n
|
||||
mean = self.mean[sample_idx]
|
||||
if len(self.mean.shape) > len(sample_idx[0].shape):
|
||||
std = self.std[:,sample_idx]+ self.eps # T*batch*n
|
||||
mean = self.mean[:,sample_idx]
|
||||
|
||||
# x is in shape of batch*n or T*batch*n
|
||||
x = (x * std) + mean
|
||||
return x
|
||||
|
||||
def cuda(self):
|
||||
self.mean = self.mean.cuda()
|
||||
self.std = self.std.cuda()
|
||||
|
||||
def cpu(self):
|
||||
self.mean = self.mean.cpu()
|
||||
self.std = self.std.cpu()
|
||||
|
||||
# normalization, Gaussian
|
||||
class GaussianNormalizer(object):
|
||||
def __init__(self, x, eps=0.00001):
|
||||
super(GaussianNormalizer, self).__init__()
|
||||
|
||||
self.mean = torch.mean(x)
|
||||
self.std = torch.std(x)
|
||||
self.eps = eps
|
||||
|
||||
def encode(self, x):
|
||||
x = (x - self.mean) / (self.std + self.eps)
|
||||
return x
|
||||
|
||||
def decode(self, x, sample_idx=None):
|
||||
x = (x * (self.std + self.eps)) + self.mean
|
||||
return x
|
||||
|
||||
def cuda(self):
|
||||
self.mean = self.mean.cuda()
|
||||
self.std = self.std.cuda()
|
||||
|
||||
def cpu(self):
|
||||
self.mean = self.mean.cpu()
|
||||
self.std = self.std.cpu()
|
||||
|
||||
|
||||
# normalization, scaling by range
|
||||
class RangeNormalizer(object):
|
||||
def __init__(self, x, low=0.0, high=1.0):
|
||||
super(RangeNormalizer, self).__init__()
|
||||
mymin = torch.min(x, 0)[0].view(-1)
|
||||
mymax = torch.max(x, 0)[0].view(-1)
|
||||
|
||||
self.a = (high - low)/(mymax - mymin)
|
||||
self.b = -self.a*mymax + high
|
||||
|
||||
def encode(self, x):
|
||||
s = x.size()
|
||||
x = x.view(s[0], -1)
|
||||
x = self.a*x + self.b
|
||||
x = x.view(s)
|
||||
return x
|
||||
|
||||
def decode(self, x):
|
||||
s = x.size()
|
||||
x = x.view(s[0], -1)
|
||||
x = (x - self.b)/self.a
|
||||
x = x.view(s)
|
||||
return x
|
||||
|
||||
class LpLoss(object):
|
||||
def __init__(self, d=2, p=2, size_average=True, reduction=True):
|
||||
super(LpLoss, self).__init__()
|
||||
|
||||
#Dimension and Lp-norm type are postive
|
||||
assert d > 0 and p > 0
|
||||
|
||||
self.d = d
|
||||
self.p = p
|
||||
self.reduction = reduction
|
||||
self.size_average = size_average
|
||||
|
||||
def abs(self, x, y):
|
||||
num_examples = x.size()[0]
|
||||
|
||||
#Assume uniform mesh
|
||||
h = 1.0 / (x.size()[1] - 1.0)
|
||||
|
||||
all_norms = (h**(self.d/self.p))*torch.norm(x.view(num_examples,-1) - y.view(num_examples,-1), self.p, 1)
|
||||
|
||||
if self.reduction:
|
||||
if self.size_average:
|
||||
return torch.mean(all_norms)
|
||||
else:
|
||||
return torch.sum(all_norms)
|
||||
|
||||
return all_norms
|
||||
|
||||
def rel(self, x, y):
|
||||
num_examples = x.size()[0]
|
||||
|
||||
diff_norms = torch.norm(x.reshape(num_examples,-1) - y.reshape(num_examples,-1), self.p, 1)
|
||||
y_norms = torch.norm(y.reshape(num_examples,-1), self.p, 1)
|
||||
|
||||
if self.reduction:
|
||||
if self.size_average:
|
||||
return torch.mean(diff_norms/y_norms)
|
||||
else:
|
||||
return torch.sum(diff_norms/y_norms)
|
||||
|
||||
return diff_norms/y_norms
|
||||
|
||||
def __call__(self, x, y):
|
||||
return self.rel(x, y)
|
||||
Vendored
-102
@@ -1,102 +0,0 @@
|
||||
# coding=utf-8
|
||||
# author=maziqing
|
||||
# email=maziqing.mzq@alibaba-inc.com
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from layers.Embed import DataEmbedding, DataEmbedding_wo_pos
|
||||
from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
|
||||
from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Autoformer is the first method to achieve the series-wise connection,
|
||||
with inherent O(LlogL) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.seq_len = configs.seq_len
|
||||
self.label_len = configs.label_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Decomp
|
||||
kernel_size = configs.moving_avg
|
||||
self.decomp = series_decomp(kernel_size)
|
||||
|
||||
# Embedding
|
||||
# The series-wise connection inherently contains the sequential information.
|
||||
# Thus, we can discard the position embedding of transformers.
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(True, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AutoCorrelationLayer(
|
||||
AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.c_out,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
# decomp init
|
||||
mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
|
||||
zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)
|
||||
seasonal_init, trend_init = self.decomp(x_enc)
|
||||
# decoder input
|
||||
trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
|
||||
seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)
|
||||
# enc
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
# dec
|
||||
dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
|
||||
seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
|
||||
trend=trend_init)
|
||||
# final
|
||||
dec_out = trend_part + seasonal_part
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
Vendored
-227
@@ -1,227 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from layers.Embed import DataEmbedding, DataEmbedding_wo_pos,DataEmbedding_wo_pos_temp,DataEmbedding_wo_temp
|
||||
from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
|
||||
from layers.FourierCorrelation import FourierBlock, FourierCrossAttention
|
||||
from layers.MultiWaveletCorrelation import MultiWaveletCross, MultiWaveletTransform
|
||||
from layers.SelfAttention_Family import FullAttention, ProbAttention
|
||||
# from layers.FED_wo_decomp import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp, series_decomp_multi
|
||||
from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp, series_decomp_multi
|
||||
import math
|
||||
import numpy as np
|
||||
|
||||
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
FEDformer performs the attention mechanism on frequency domain and achieved O(N) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.version = configs.version
|
||||
self.mode_select = configs.mode_select
|
||||
self.modes = configs.modes
|
||||
self.seq_len = configs.seq_len
|
||||
self.label_len = configs.label_len
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Decomp
|
||||
kernel_size = configs.moving_avg
|
||||
if isinstance(kernel_size, list):
|
||||
self.decomp = series_decomp_multi(kernel_size)
|
||||
else:
|
||||
self.decomp = series_decomp(kernel_size)
|
||||
|
||||
# Embedding
|
||||
# The series-wise connection inherently contains the sequential information.
|
||||
# Thus, we can discard the position embedding of transformers.
|
||||
# self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
# configs.dropout)
|
||||
# self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
# configs.dropout)
|
||||
if configs.embed_type == 0:
|
||||
self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 1:
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 2:
|
||||
self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
elif configs.embed_type == 3:
|
||||
self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
if configs.version == 'Wavelets':
|
||||
encoder_self_att = MultiWaveletTransform(ich=configs.d_model, L=configs.L, base=configs.base)
|
||||
decoder_self_att = MultiWaveletTransform(ich=configs.d_model, L=configs.L, base=configs.base)
|
||||
decoder_cross_att = MultiWaveletCross(in_channels=configs.d_model,
|
||||
out_channels=configs.d_model,
|
||||
seq_len_q=self.seq_len // 2 + self.pred_len,
|
||||
seq_len_kv=self.seq_len,
|
||||
modes=configs.modes,
|
||||
ich=configs.d_model,
|
||||
base=configs.base,
|
||||
activation=configs.cross_activation)
|
||||
else:
|
||||
encoder_self_att = FourierBlock(in_channels=configs.d_model,
|
||||
out_channels=configs.d_model,
|
||||
seq_len=self.seq_len,
|
||||
modes=configs.modes,
|
||||
mode_select_method=configs.mode_select)
|
||||
decoder_self_att = FourierBlock(in_channels=configs.d_model,
|
||||
out_channels=configs.d_model,
|
||||
seq_len=self.seq_len//2+self.pred_len,
|
||||
modes=configs.modes,
|
||||
mode_select_method=configs.mode_select)
|
||||
decoder_cross_att = FourierCrossAttention(in_channels=configs.d_model,
|
||||
out_channels=configs.d_model,
|
||||
seq_len_q=self.seq_len//2+self.pred_len,
|
||||
seq_len_kv=self.seq_len,
|
||||
modes=configs.modes,
|
||||
mode_select_method=configs.mode_select)
|
||||
# Encoder
|
||||
enc_modes = int(min(configs.modes, configs.seq_len//2))
|
||||
dec_modes = int(min(configs.modes, (configs.seq_len//2+configs.pred_len)//2))
|
||||
print('enc_modes: {}, dec_modes: {}'.format(enc_modes, dec_modes))
|
||||
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
encoder_self_att,
|
||||
configs.d_model, configs.n_heads),
|
||||
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AutoCorrelationLayer(
|
||||
decoder_self_att,
|
||||
configs.d_model, configs.n_heads),
|
||||
AutoCorrelationLayer(
|
||||
decoder_cross_att,
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.c_out,
|
||||
configs.d_ff,
|
||||
moving_avg=configs.moving_avg,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=my_Layernorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
# def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
# enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
# # decomp init
|
||||
# mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
|
||||
# # zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]]).to(device) # cuda()
|
||||
# # seasonal_init, trend_init = self.decomp(x_enc)
|
||||
# # decoder input
|
||||
# dec_in = torch.cat([x_enc[:, -self.label_len:, :], mean], dim=1)
|
||||
# dec_out = F.pad(x_enc[:, -self.label_len:, :], (0, 0, 0, self.pred_len))
|
||||
# # enc
|
||||
# enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
# enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
# # dec
|
||||
# dec_out = self.dec_embedding(dec_out, x_mark_dec)
|
||||
# dec_out, _ = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
|
||||
# trend=dec_in)
|
||||
# # final
|
||||
# # dec_out = trend_part + seasonal_part
|
||||
|
||||
# if self.output_attention:
|
||||
# return dec_out[:, -self.pred_len:, :], attns
|
||||
# else:
|
||||
# return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
# decomp init
|
||||
mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
|
||||
zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]]).to(device) # cuda()
|
||||
seasonal_init, trend_init = self.decomp(x_enc)
|
||||
# decoder input
|
||||
trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
|
||||
seasonal_init = F.pad(seasonal_init[:, -self.label_len:, :], (0, 0, 0, self.pred_len))
|
||||
# enc
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
# dec
|
||||
dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
|
||||
seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
|
||||
trend=trend_init)
|
||||
# final
|
||||
dec_out = trend_part + seasonal_part
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
|
||||
if __name__ == '__main__':
|
||||
class Configs(object):
|
||||
ab = 0
|
||||
modes = 32
|
||||
mode_select = 'random'
|
||||
# version = 'Fourier'
|
||||
version = 'Wavelets'
|
||||
moving_avg = [12, 24]
|
||||
L = 1
|
||||
base = 'legendre'
|
||||
cross_activation = 'tanh'
|
||||
seq_len = 96
|
||||
label_len = 48
|
||||
pred_len = 96
|
||||
output_attention = True
|
||||
enc_in = 7
|
||||
dec_in = 7
|
||||
d_model = 16
|
||||
embed = 'timeF'
|
||||
dropout = 0.05
|
||||
freq = 'h'
|
||||
factor = 1
|
||||
n_heads = 8
|
||||
d_ff = 16
|
||||
e_layers = 2
|
||||
d_layers = 1
|
||||
c_out = 7
|
||||
activation = 'gelu'
|
||||
wavelet = 0
|
||||
|
||||
configs = Configs()
|
||||
model = Model(configs)
|
||||
|
||||
print('parameter number is {}'.format(sum(p.numel() for p in model.parameters())))
|
||||
enc = torch.randn([3, configs.seq_len, 7])
|
||||
enc_mark = torch.randn([3, configs.seq_len, 4])
|
||||
|
||||
dec = torch.randn([3, configs.seq_len//2+configs.pred_len, 7])
|
||||
dec_mark = torch.randn([3, configs.seq_len//2+configs.pred_len, 4])
|
||||
out = model.forward(enc, enc_mark, dec, dec_mark)
|
||||
print(out)
|
||||
Vendored
-80
@@ -1,80 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from utils.masking import TriangularCausalMask, ProbMask
|
||||
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
|
||||
from layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer
|
||||
from layers.Embed import DataEmbedding
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Informer with Propspare attention in O(LlogL) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Embedding
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AttentionLayer(
|
||||
ProbAttention(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
[
|
||||
ConvLayer(
|
||||
configs.d_model
|
||||
) for l in range(configs.e_layers - 1)
|
||||
] if configs.distil else None,
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AttentionLayer(
|
||||
ProbAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AttentionLayer(
|
||||
ProbAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
|
||||
dec_out = self.dec_embedding(x_dec, x_mark_dec)
|
||||
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
Vendored
-71
@@ -1,71 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
|
||||
from layers.SelfAttention_Family import FullAttention, AttentionLayer
|
||||
from layers.Embed import DataEmbedding
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
Vanilla Transformer with O(L^2) complexity
|
||||
"""
|
||||
def __init__(self, configs):
|
||||
super(Model, self).__init__()
|
||||
self.pred_len = configs.pred_len
|
||||
self.output_attention = configs.output_attention
|
||||
|
||||
# Embedding
|
||||
self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
|
||||
configs.dropout)
|
||||
# Encoder
|
||||
self.encoder = Encoder(
|
||||
[
|
||||
EncoderLayer(
|
||||
AttentionLayer(
|
||||
FullAttention(False, configs.factor, attention_dropout=configs.dropout,
|
||||
output_attention=configs.output_attention), configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation
|
||||
) for l in range(configs.e_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model)
|
||||
)
|
||||
# Decoder
|
||||
self.decoder = Decoder(
|
||||
[
|
||||
DecoderLayer(
|
||||
AttentionLayer(
|
||||
FullAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
AttentionLayer(
|
||||
FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
|
||||
configs.d_model, configs.n_heads),
|
||||
configs.d_model,
|
||||
configs.d_ff,
|
||||
dropout=configs.dropout,
|
||||
activation=configs.activation,
|
||||
)
|
||||
for l in range(configs.d_layers)
|
||||
],
|
||||
norm_layer=torch.nn.LayerNorm(configs.d_model),
|
||||
projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
|
||||
enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
|
||||
|
||||
enc_out = self.enc_embedding(x_enc, x_mark_enc)
|
||||
enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
|
||||
|
||||
dec_out = self.dec_embedding(x_dec, x_mark_dec)
|
||||
dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
|
||||
|
||||
if self.output_attention:
|
||||
return dec_out[:, -self.pred_len:, :], attns
|
||||
else:
|
||||
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
||||
-165
@@ -1,165 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
import torch
|
||||
from exp.exp_main import Exp_Main
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
fix_seed = 2021
|
||||
random.seed(fix_seed)
|
||||
torch.manual_seed(fix_seed)
|
||||
np.random.seed(fix_seed)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Autoformer & Transformer family for Time Series Forecasting')
|
||||
|
||||
# basic config
|
||||
parser.add_argument('--is_training', type=int, default=1, help='status')
|
||||
parser.add_argument('--task_id', type=str, default='test', help='task id')
|
||||
parser.add_argument('--model', type=str, default='Reformer',
|
||||
help='model name, options: [FEDformer, Autoformer, Informer, Transformer]')
|
||||
|
||||
# supplementary config for FEDformer model
|
||||
parser.add_argument('--version', type=str, default='Fourier',
|
||||
help='for FEDformer, there are two versions to choose, options: [Fourier, Wavelets]')
|
||||
parser.add_argument('--mode_select', type=str, default='random',
|
||||
help='for FEDformer, there are two mode selection method, options: [random, low]')
|
||||
parser.add_argument('--modes', type=int, default=64, help='modes to be selected random 64')
|
||||
parser.add_argument('--L', type=int, default=3, help='ignore level')
|
||||
parser.add_argument('--base', type=str, default='legendre', help='mwt base')
|
||||
parser.add_argument('--cross_activation', type=str, default='tanh',
|
||||
help='mwt cross atention activation function tanh or softmax')
|
||||
|
||||
|
||||
# data loader
|
||||
parser.add_argument('--data', type=str, default='ETTh1', help='dataset type')
|
||||
parser.add_argument('--root_path', type=str, default='../dataset', help='root path of the data file')
|
||||
parser.add_argument('--data_path', type=str, default='ETTh1.csv', help='data file')
|
||||
parser.add_argument('--features', type=str, default='M',
|
||||
help='forecasting task, options:[M, S, MS]; M:multivariate predict multivariate, '
|
||||
'S:univariate predict univariate, MS:multivariate predict univariate')
|
||||
parser.add_argument('--target', type=str, default='OT', help='target feature in S or MS task')
|
||||
parser.add_argument('--freq', type=str, default='h',
|
||||
help='freq for time features encoding, options:[s:secondly, t:minutely, h:hourly, d:daily, '
|
||||
'b:business days, w:weekly, m:monthly], you can also use more detailed freq like 15min or 3h')
|
||||
parser.add_argument('--checkpoints', type=str, default='./checkpoints/', help='location of model checkpoints')
|
||||
|
||||
# forecasting task
|
||||
parser.add_argument('--seq_len', type=int, default=96, help='input sequence length')
|
||||
parser.add_argument('--label_len', type=int, default=48, help='start token length')
|
||||
parser.add_argument('--pred_len', type=int, default=96, help='prediction sequence length')
|
||||
parser.add_argument('--embed_type', type=int, default=0, help='prediction sequence length')
|
||||
# parser.add_argument('--cross_activation', type=str, default='tanh'
|
||||
|
||||
# model define
|
||||
parser.add_argument('--enc_in', type=int, default=7, help='encoder input size')
|
||||
parser.add_argument('--dec_in', type=int, default=7, help='decoder input size')
|
||||
parser.add_argument('--c_out', type=int, default=7, help='output size')
|
||||
parser.add_argument('--d_model', type=int, default=512, help='dimension of model')
|
||||
parser.add_argument('--n_heads', type=int, default=8, help='num of heads')
|
||||
parser.add_argument('--e_layers', type=int, default=2, help='num of encoder layers')
|
||||
parser.add_argument('--d_layers', type=int, default=1, help='num of decoder layers')
|
||||
parser.add_argument('--d_ff', type=int, default=2048, help='dimension of fcn')
|
||||
parser.add_argument('--moving_avg', default=[24], help='window size of moving average')
|
||||
parser.add_argument('--factor', type=int, default=1, help='attn factor')
|
||||
parser.add_argument('--distil', action='store_false',
|
||||
help='whether to use distilling in encoder, using this argument means not using distilling',
|
||||
default=True)
|
||||
parser.add_argument('--dropout', type=float, default=0.05, help='dropout')
|
||||
parser.add_argument('--embed', type=str, default='timeF',
|
||||
help='time features encoding, options:[timeF, fixed, learned]')
|
||||
parser.add_argument('--activation', type=str, default='gelu', help='activation')
|
||||
parser.add_argument('--output_attention', action='store_true', help='whether to output attention in ecoder')
|
||||
parser.add_argument('--do_predict', action='store_true', help='whether to predict unseen future data')
|
||||
|
||||
# optimization
|
||||
parser.add_argument('--num_workers', type=int, default=10, help='data loader num workers')
|
||||
parser.add_argument('--itr', type=int, default=1, help='experiments times')
|
||||
parser.add_argument('--train_epochs', type=int, default=10, help='train epochs')
|
||||
parser.add_argument('--batch_size', type=int, default=32, help='batch size of train input data')
|
||||
parser.add_argument('--patience', type=int, default=3, help='early stopping patience')
|
||||
parser.add_argument('--learning_rate', type=float, default=0.0001, help='optimizer learning rate')
|
||||
parser.add_argument('--des', type=str, default='test', help='exp description')
|
||||
parser.add_argument('--loss', type=str, default='mse', help='loss function')
|
||||
parser.add_argument('--lradj', type=str, default='type1', help='adjust learning rate')
|
||||
parser.add_argument('--use_amp', action='store_true', help='use automatic mixed precision training', default=False)
|
||||
|
||||
# GPU
|
||||
parser.add_argument('--use_gpu', type=bool, default=True, help='use gpu')
|
||||
parser.add_argument('--gpu', type=int, default=0, help='gpu')
|
||||
parser.add_argument('--use_multi_gpu', action='store_true', help='use multiple gpus', default=False)
|
||||
parser.add_argument('--devices', type=str, default='0,1', help='device ids of multi gpus')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
args.use_gpu = True if torch.cuda.is_available() and args.use_gpu else False
|
||||
|
||||
if args.use_gpu and args.use_multi_gpu:
|
||||
args.dvices = args.devices.replace(' ', '')
|
||||
device_ids = args.devices.split(',')
|
||||
args.device_ids = [int(id_) for id_ in device_ids]
|
||||
args.gpu = args.device_ids[0]
|
||||
|
||||
print('Args in experiment:')
|
||||
print(args)
|
||||
|
||||
Exp = Exp_Main
|
||||
|
||||
if args.is_training:
|
||||
for ii in range(args.itr):
|
||||
# setting record of experiments
|
||||
setting = '{}_{}_{}_modes{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(
|
||||
args.task_id,
|
||||
args.model,
|
||||
args.mode_select,
|
||||
args.modes,
|
||||
args.data,
|
||||
args.features,
|
||||
args.seq_len,
|
||||
args.label_len,
|
||||
args.pred_len,
|
||||
args.d_model,
|
||||
args.n_heads,
|
||||
args.e_layers,
|
||||
args.d_layers,
|
||||
args.d_ff,
|
||||
args.factor,
|
||||
args.embed,
|
||||
args.distil,
|
||||
args.des,
|
||||
ii)
|
||||
|
||||
exp = Exp(args) # set experiments
|
||||
print('>>>>>>>start training : {}>>>>>>>>>>>>>>>>>>>>>>>>>>'.format(setting))
|
||||
exp.train(setting)
|
||||
|
||||
print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
|
||||
exp.test(setting)
|
||||
|
||||
if args.do_predict:
|
||||
print('>>>>>>>predicting : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
|
||||
exp.predict(setting, True)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
else:
|
||||
ii = 0
|
||||
setting = '{}_{}_{}_ft{}_sl{}_ll{}_pl{}_dm{}_nh{}_el{}_dl{}_df{}_fc{}_eb{}_dt{}_{}_{}'.format(args.model_id,
|
||||
args.model,
|
||||
args.data,
|
||||
args.features,
|
||||
args.seq_len,
|
||||
args.label_len,
|
||||
args.pred_len,
|
||||
args.d_model,
|
||||
args.n_heads,
|
||||
args.e_layers,
|
||||
args.d_layers,
|
||||
args.d_ff,
|
||||
args.factor,
|
||||
args.embed,
|
||||
args.distil,
|
||||
args.des, ii)
|
||||
|
||||
exp = Exp(args) # set experiments
|
||||
print('>>>>>>>testing : {}<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<'.format(setting))
|
||||
exp.test(setting, test=1)
|
||||
torch.cuda.empty_cache()
|
||||
-202
@@ -1,202 +0,0 @@
|
||||
# cd FEDformer
|
||||
if [ ! -d "../logs" ]; then
|
||||
mkdir ../logs
|
||||
fi
|
||||
|
||||
if [ ! -d "../logs/LongForecasting" ]; then
|
||||
mkdir ../logs/LongForecasting
|
||||
fi
|
||||
|
||||
for preLen in 96 192 336 720
|
||||
do
|
||||
# ETTm1
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path ETTm1.csv \
|
||||
--task_id ETTm1 \
|
||||
--model FEDformer \
|
||||
--data ETTm1 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_ETTm1_$pred_len.log
|
||||
|
||||
# ETTh1
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path ETTh1.csv \
|
||||
--task_id ETTh1 \
|
||||
--model FEDformer \
|
||||
--data ETTh1 \
|
||||
--features S \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_ETTh1_$pred_len.log
|
||||
|
||||
# ETTm2
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path ETTm2.csv \
|
||||
--task_id ETTm2 \
|
||||
--model FEDformer \
|
||||
--data ETTm2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_ETTm2_$pred_len.log
|
||||
|
||||
# ETTh2
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path ETTh2.csv \
|
||||
--task_id ETTh2 \
|
||||
--model FEDformer \
|
||||
--data ETTh2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_ETTh2_$pred_len.log
|
||||
|
||||
# electricity
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path electricity.csv \
|
||||
--task_id ECL \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 321 \
|
||||
--dec_in 321 \
|
||||
--c_out 321 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_electricity_$pred_len.log
|
||||
|
||||
# exchange
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path exchange_rate.csv \
|
||||
--task_id Exchange \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features S \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 8 \
|
||||
--dec_in 8 \
|
||||
--c_out 8 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_exchange_rate_$pred_len.log
|
||||
|
||||
# traffic
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path traffic.csv \
|
||||
--task_id traffic \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features S \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 862 \
|
||||
--dec_in 862 \
|
||||
--c_out 862 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 3 >../logs/LongForecasting/FEDformer_traffic_$pred_len.log
|
||||
|
||||
# weather
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path weather.csv \
|
||||
--task_id weather \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features S \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 21 \
|
||||
--dec_in 21 \
|
||||
--c_out 21 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_weather_$pred_len.log
|
||||
done
|
||||
|
||||
|
||||
for preLen in 24 36 48 60
|
||||
do
|
||||
# illness
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--data_path national_illness.csv \
|
||||
--task_id ili \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features S \
|
||||
--seq_len 36 \
|
||||
--label_len 18 \
|
||||
--pred_len $preLen \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LongForecasting/FEDformer_ili_$pred_len.log
|
||||
done
|
||||
|
||||
# cd ..
|
||||
-237
@@ -1,237 +0,0 @@
|
||||
# cd FEDformer
|
||||
if [ ! -d "../logs" ]; then
|
||||
mkdir ../logs
|
||||
fi
|
||||
|
||||
if [ ! -d "../logs/LookBackWindow" ]; then
|
||||
mkdir ../logs/LookBackWindow
|
||||
fi
|
||||
|
||||
for seqLen in 36 48 60 72 144 288
|
||||
do
|
||||
for pred_len in 24 576
|
||||
do
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path ETTm1.csv \
|
||||
--task_id ETTm1 \
|
||||
--model FEDformer \
|
||||
--data ETTm1 \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ETTm2_$seqLen'_'$pred_len.log
|
||||
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path ETTm2.csv \
|
||||
--task_id ETTm2 \
|
||||
--model FEDformer \
|
||||
--data ETTm2 \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ETTm2_$seqLen'_'$pred_len.log
|
||||
done
|
||||
done
|
||||
|
||||
for seqLen in 48 72 120 144 168 192 336 720
|
||||
do
|
||||
for pred_len in 24 720
|
||||
do
|
||||
# ETTh1
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path ETTh1.csv \
|
||||
--task_id ETTh1 \
|
||||
--model FEDformer \
|
||||
--data ETTh1 \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ETTh1_$seqLen'_'$pred_len.log
|
||||
|
||||
# ETTh2
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path ETTh2.csv \
|
||||
--task_id ETTh2 \
|
||||
--model FEDformer \
|
||||
--data ETTh2 \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--d_model 512 \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ETTh2_$seqLen'_'$pred_len.log
|
||||
|
||||
## electricity
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path electricity.csv \
|
||||
--task_id ECL \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 321 \
|
||||
--dec_in 321 \
|
||||
--c_out 321 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_electricity_$seqLen'_'$pred_len.log
|
||||
|
||||
# exchange
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path exchange_rate.csv \
|
||||
--task_id Exchange \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 8 \
|
||||
--dec_in 8 \
|
||||
--c_out 8 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_exchange_rate_$seqLen'_'$pred_len.log
|
||||
|
||||
# traffic
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path traffic.csv \
|
||||
--task_id traffic \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 862 \
|
||||
--dec_in 862 \
|
||||
--c_out 862 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 3 >../logs/LookBackWindow/FEDformer_traffic_$seqLen'_'$pred_len.log
|
||||
|
||||
# weather
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path weather.csv \
|
||||
--task_id weather \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 21 \
|
||||
--dec_in 21 \
|
||||
--c_out 21 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_weather_$seqLen'_'$pred_len.log
|
||||
done
|
||||
done
|
||||
|
||||
|
||||
for seqLen in 26 52 78 104 130 156 208
|
||||
do
|
||||
# illness
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path national_illness.csv \
|
||||
--task_id ili \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 18 \
|
||||
--pred_len 24 \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ili_$seqLen'_'24.log
|
||||
|
||||
python -u run.py \
|
||||
--is_training 1 \
|
||||
--root_path .../dataset/ \
|
||||
--data_path national_illness.csv \
|
||||
--task_id ili \
|
||||
--model FEDformer \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len $seqLen \
|
||||
--label_len 18 \
|
||||
--pred_len 60 \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >../logs/LookBackWindow/FEDformer_ili_$seqLen'_'60.log
|
||||
done
|
||||
# cd ..
|
||||
Vendored
-39
@@ -1,39 +0,0 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
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
|
||||
|
||||
class LocalMask():
|
||||
def __init__(self, B, L,S,device="cpu"):
|
||||
mask_shape = [B, 1, L, S]
|
||||
with torch.no_grad():
|
||||
self.len = math.ceil(np.log2(L))
|
||||
self._mask1 = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
|
||||
self._mask2 = ~torch.triu(torch.ones(mask_shape,dtype=torch.bool),diagonal=-self.len).to(device)
|
||||
self._mask = self._mask1+self._mask2
|
||||
@property
|
||||
def mask(self):
|
||||
return self._mask
|
||||
Vendored
-51
@@ -1,51 +0,0 @@
|
||||
import numpy as np
|
||||
|
||||
|
||||
def RSE(pred, true):
|
||||
return np.sqrt(np.sum((true - pred) ** 2)) / np.sqrt(np.sum((true - true.mean()) ** 2))
|
||||
|
||||
|
||||
def CORR(pred, true):
|
||||
u = ((true - true.mean(0)) * (pred - pred.mean(0))).sum(0)
|
||||
d = np.sqrt(((true - true.mean(0)) ** 2 * (pred - pred.mean(0)) ** 2).sum(0))
|
||||
return (u / d).mean(-1)
|
||||
|
||||
|
||||
def MAE(pred, true):
|
||||
return np.mean(np.abs(pred - true))
|
||||
|
||||
|
||||
def MSE(pred, true):
|
||||
return np.mean((pred - true) ** 2)
|
||||
|
||||
|
||||
def RMSE(pred, true):
|
||||
return np.sqrt(MSE(pred, true))
|
||||
|
||||
|
||||
def MAPE(pred, true):
|
||||
return np.mean(np.abs((pred - true) / true))
|
||||
|
||||
|
||||
def MSPE(pred, true):
|
||||
return np.mean(np.square((pred - true) / true))
|
||||
|
||||
|
||||
def metric(pred, true):
|
||||
mae = MAE(pred, true)
|
||||
mse = MSE(pred, true)
|
||||
rmse = RMSE(pred, true)
|
||||
mape = MAPE(pred, true)
|
||||
mspe = MSPE(pred, true)
|
||||
|
||||
return mae, mse, rmse, mape, mspe
|
||||
|
||||
def metric2(pred, true):
|
||||
mae = MAE(pred, true)
|
||||
mse = MSE(pred, true)
|
||||
rmse = RMSE(pred, true)
|
||||
mape = MAPE(pred, true)
|
||||
mspe = MSPE(pred, true)
|
||||
rse = RSE(pred, true)
|
||||
corr = CORR(pred, true)
|
||||
return mae, mse, rmse, mape, mspe, rse, corr
|
||||
Vendored
-134
@@ -1,134 +0,0 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.tseries import offsets
|
||||
from pandas.tseries.frequencies import to_offset
|
||||
|
||||
|
||||
class TimeFeature:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
pass
|
||||
|
||||
def __repr__(self):
|
||||
return self.__class__.__name__ + "()"
|
||||
|
||||
|
||||
class SecondOfMinute(TimeFeature):
|
||||
"""Minute of hour encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.second / 59.0 - 0.5
|
||||
|
||||
|
||||
class MinuteOfHour(TimeFeature):
|
||||
"""Minute of hour encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.minute / 59.0 - 0.5
|
||||
|
||||
|
||||
class HourOfDay(TimeFeature):
|
||||
"""Hour of day encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.hour / 23.0 - 0.5
|
||||
|
||||
|
||||
class DayOfWeek(TimeFeature):
|
||||
"""Hour of day encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return index.dayofweek / 6.0 - 0.5
|
||||
|
||||
|
||||
class DayOfMonth(TimeFeature):
|
||||
"""Day of month encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.day - 1) / 30.0 - 0.5
|
||||
|
||||
|
||||
class DayOfYear(TimeFeature):
|
||||
"""Day of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.dayofyear - 1) / 365.0 - 0.5
|
||||
|
||||
|
||||
class MonthOfYear(TimeFeature):
|
||||
"""Month of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.month - 1) / 11.0 - 0.5
|
||||
|
||||
|
||||
class WeekOfYear(TimeFeature):
|
||||
"""Week of year encoded as value between [-0.5, 0.5]"""
|
||||
|
||||
def __call__(self, index: pd.DatetimeIndex) -> np.ndarray:
|
||||
return (index.isocalendar().week - 1) / 52.0 - 0.5
|
||||
|
||||
|
||||
def time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]:
|
||||
"""
|
||||
Returns a list of time features that will be appropriate for the given frequency string.
|
||||
Parameters
|
||||
----------
|
||||
freq_str
|
||||
Frequency string of the form [multiple][granularity] such as "12H", "5min", "1D" etc.
|
||||
"""
|
||||
|
||||
features_by_offsets = {
|
||||
offsets.YearEnd: [],
|
||||
offsets.QuarterEnd: [MonthOfYear],
|
||||
offsets.MonthEnd: [MonthOfYear],
|
||||
offsets.Week: [DayOfMonth, WeekOfYear],
|
||||
offsets.Day: [DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.BusinessDay: [DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.Hour: [HourOfDay, DayOfWeek, DayOfMonth, DayOfYear],
|
||||
offsets.Minute: [
|
||||
MinuteOfHour,
|
||||
HourOfDay,
|
||||
DayOfWeek,
|
||||
DayOfMonth,
|
||||
DayOfYear,
|
||||
],
|
||||
offsets.Second: [
|
||||
SecondOfMinute,
|
||||
MinuteOfHour,
|
||||
HourOfDay,
|
||||
DayOfWeek,
|
||||
DayOfMonth,
|
||||
DayOfYear,
|
||||
],
|
||||
}
|
||||
|
||||
offset = to_offset(freq_str)
|
||||
|
||||
for offset_type, feature_classes in features_by_offsets.items():
|
||||
if isinstance(offset, offset_type):
|
||||
return [cls() for cls in feature_classes]
|
||||
|
||||
supported_freq_msg = f"""
|
||||
Unsupported frequency {freq_str}
|
||||
The following frequencies are supported:
|
||||
Y - yearly
|
||||
alias: A
|
||||
M - monthly
|
||||
W - weekly
|
||||
D - daily
|
||||
B - business days
|
||||
H - hourly
|
||||
T - minutely
|
||||
alias: min
|
||||
S - secondly
|
||||
"""
|
||||
raise RuntimeError(supported_freq_msg)
|
||||
|
||||
|
||||
def time_features(dates, freq='h'):
|
||||
return np.vstack([feat(dates) for feat in time_features_from_frequency_str(freq)])
|
||||
-88
@@ -1,88 +0,0 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
plt.switch_backend('agg')
|
||||
|
||||
|
||||
def adjust_learning_rate(optimizer, epoch, args):
|
||||
# lr = args.learning_rate * (0.2 ** (epoch // 2))
|
||||
if args.lradj == 'type1':
|
||||
lr_adjust = {epoch: args.learning_rate * (0.5 ** ((epoch - 1) // 1))}
|
||||
elif args.lradj == 'type2':
|
||||
lr_adjust = {
|
||||
2: 5e-5, 4: 1e-5, 6: 5e-6, 8: 1e-6,
|
||||
10: 5e-7, 15: 1e-7, 20: 5e-8
|
||||
}
|
||||
elif args.lradj =='type3':
|
||||
lr_adjust = {epoch: args.learning_rate}
|
||||
elif args.lradj == 'type4':
|
||||
lr_adjust = {epoch: args.learning_rate * (0.9 ** ((epoch - 1) // 1))}
|
||||
if epoch in lr_adjust.keys():
|
||||
lr = lr_adjust[epoch]
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
print('Updating learning rate to {}'.format(lr))
|
||||
|
||||
|
||||
class EarlyStopping:
|
||||
def __init__(self, patience=7, verbose=False, delta=0):
|
||||
self.patience = patience
|
||||
self.verbose = verbose
|
||||
self.counter = 0
|
||||
self.best_score = None
|
||||
self.early_stop = False
|
||||
self.val_loss_min = np.Inf
|
||||
self.delta = delta
|
||||
|
||||
def __call__(self, val_loss, model, path):
|
||||
score = -val_loss
|
||||
if self.best_score is None:
|
||||
self.best_score = score
|
||||
self.save_checkpoint(val_loss, model, path)
|
||||
elif score < self.best_score + self.delta:
|
||||
self.counter += 1
|
||||
print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
|
||||
if self.counter >= self.patience:
|
||||
self.early_stop = True
|
||||
else:
|
||||
self.best_score = score
|
||||
self.save_checkpoint(val_loss, model, path)
|
||||
self.counter = 0
|
||||
|
||||
def save_checkpoint(self, val_loss, model, path):
|
||||
if self.verbose:
|
||||
print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...')
|
||||
torch.save(model.state_dict(), path + '/' + 'checkpoint.pth')
|
||||
self.val_loss_min = val_loss
|
||||
|
||||
|
||||
class dotdict(dict):
|
||||
"""dot.notation access to dictionary attributes"""
|
||||
__getattr__ = dict.get
|
||||
__setattr__ = dict.__setitem__
|
||||
__delattr__ = dict.__delitem__
|
||||
|
||||
|
||||
class StandardScaler():
|
||||
def __init__(self, mean, std):
|
||||
self.mean = mean
|
||||
self.std = std
|
||||
|
||||
def transform(self, data):
|
||||
return (data - self.mean) / self.std
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return (data * self.std) + self.mean
|
||||
|
||||
|
||||
def visual(true, preds=None, name='./pic/test.pdf'):
|
||||
"""
|
||||
Results visualization
|
||||
"""
|
||||
plt.figure()
|
||||
plt.plot(true, label='GroundTruth', linewidth=2)
|
||||
if preds is not None:
|
||||
plt.plot(preds, label='Prediction', linewidth=2)
|
||||
plt.legend()
|
||||
plt.savefig(name, bbox_inches='tight')
|
||||
-209
@@ -1,209 +0,0 @@
|
||||
# ALL scripts in this file come from Autoformer
|
||||
if [ ! -d "./logs" ]; then
|
||||
mkdir ./logs
|
||||
fi
|
||||
|
||||
if [ ! -d "./logs/LongForecasting" ]; then
|
||||
mkdir ./logs/LongForecasting
|
||||
fi
|
||||
|
||||
random_seed=2021
|
||||
model_name=Informer
|
||||
|
||||
for pred_len in 96 192 336 720
|
||||
do
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path exchange_rate.csv \
|
||||
--model_id exchange_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 8 \
|
||||
--dec_in 8 \
|
||||
--c_out 8 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 1 >logs/LongForecasting/$model_name'_exchange_rate_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path electricity.csv \
|
||||
--model_id electricity_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 321 \
|
||||
--dec_in 321 \
|
||||
--c_out 321 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_electricity_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path traffic.csv \
|
||||
--model_id traffic_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 862 \
|
||||
--dec_in 862 \
|
||||
--c_out 862 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 3 >logs/LongForecasting/$model_name'_traffic_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path weather.csv \
|
||||
--model_id weather_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 21 \
|
||||
--dec_in 21 \
|
||||
--c_out 21 \
|
||||
--des 'Exp' \
|
||||
--itr 1 \
|
||||
--train_epochs 2 >logs/LongForecasting/$model_name'_weather_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTh1.csv \
|
||||
--model_id ETTh1_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTh1 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Etth1_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTh2.csv \
|
||||
--model_id ETTh2_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTh2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Etth2_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTm1.csv \
|
||||
--model_id ETTm1_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTm1 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Ettm1_'$pred_len.log
|
||||
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path ETTm2.csv \
|
||||
--model_id ETTm2_96_$pred_len \
|
||||
--model $model_name \
|
||||
--data ETTm2 \
|
||||
--features M \
|
||||
--seq_len 96 \
|
||||
--label_len 48 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_Ettm2_'$pred_len.log
|
||||
done
|
||||
|
||||
for pred_len in 24 36 48 60
|
||||
do
|
||||
python -u run_longExp.py \
|
||||
--random_seed $random_seed \
|
||||
--is_training 1 \
|
||||
--root_path ./dataset/ \
|
||||
--data_path national_illness.csv \
|
||||
--model_id ili_36_$pred_len \
|
||||
--model $model_name \
|
||||
--data custom \
|
||||
--features M \
|
||||
--seq_len 36 \
|
||||
--label_len 18 \
|
||||
--pred_len $pred_len \
|
||||
--e_layers 2 \
|
||||
--d_layers 1 \
|
||||
--factor 3 \
|
||||
--enc_in 7 \
|
||||
--dec_in 7 \
|
||||
--c_out 7 \
|
||||
--des 'Exp' \
|
||||
--itr 1 >logs/LongForecasting/$model_name'_ili_'$pred_len.log
|
||||
done
|
||||
-7
@@ -1,7 +0,0 @@
|
||||
Legal Disclaimer
|
||||
|
||||
Within this source code, the comments in Chinese shall be the original, governing version. Any comment in other languages are for reference only. In the event of any conflict between the Chinese language version comments and other language version comments, the Chinese language version shall prevail.
|
||||
|
||||
法律免责声明
|
||||
|
||||
关于代码注释部分,中文注释为官方版本,其它语言注释仅做参考。中文注释可能与其它语言注释存在不一致,当中文注释与其它语言注释存在不一致时,请以中文注释为准。
|
||||
-203
@@ -1,203 +0,0 @@
|
||||
Copyright 2022 - The AntGroup-RI team. All rights reserved.
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
-112
@@ -1,112 +0,0 @@
|
||||
# Pyraformer: Low-complexity Pyramidal Attention for Long-range Time Series Modeling and Forecasting
|
||||
This is the Pytorch implementation of Pyraformer (Pyramidal Attention based Transformer) in the ICLR paper: [Pyraformer: Low-complexity Pyramidal Attention for Long-range Time Series Modeling and Forecasting](https://openreview.net/pdf?id=0EXmFzUn5I).
|
||||
From https://github.com/alipay/Pyraformer
|
||||

|
||||
<center><b>Figure 1.</b> The network architecture of Pyraformer.</center>
|
||||
|
||||
## Pyramidal Attention
|
||||
As demonstrated in Figure 2, we leverage a pyramidal graph to describe the temporal dependencies of the observed time series in a multiresolution fashion. We can decompose the pyramidal graph into two parts: the inter-scale and the intra-scale connections. The inter-scale connections form a C-ary tree, in which each parent has C children. For example, if we associate the finest scale of the pyramidal graph with hourly observations of the original time series, the nodes at coarser scales can be regarded as the daily, weekly, and even monthly features of the time series. As a consequence, the pyramidal graph offers a multiresolution representation of the original time series. Furthermore, it is easier to capture long-range dependencies (e.g., monthly dependence) in the coarser scales by simply connecting the neighboring nodes via the intra-scale connections. In other words, the coarser scales are instrumental in describing long-range correlations in a manner that is graphically far more parsimonious than could be solely captured with a single, finest scale model.
|
||||
|
||||
|
||||

|
||||
<center><b>Figure 2.</b> The Pyramidal Attention Mechanism.</center>
|
||||
|
||||
## Requirements
|
||||
* Python 3.7
|
||||
* pytorch 1.8.0
|
||||
* CUDA 11.1
|
||||
* TVM 0.8.0 (optional)
|
||||
|
||||
Dependencies can be installed by:
|
||||
|
||||
pip install -r requirements.txt
|
||||
|
||||
If you are using CUDA 11.1, you can use the compiled TVM runtime version in the our code to run PAM-TVM. Due to the short history length in the experiments, PAM-TVM does not provide a speed increase. If you want to compile our PAM-TVM kernel yourself, see [here](https://tvm.apache.org/docs/install/index.html) to compile TVM 0.8.0 first.
|
||||
|
||||
## Data preparetion
|
||||
The four datasets (Electricity, Wind, ETT and App Flow) used in this paper can be downloaded from the following links:
|
||||
* [Electricity](https://archive.ics.uci.edu/ml/datasets/ElectricityLoadDiagrams20112014)
|
||||
* [Wind](https://www.kaggle.com/sohier/30-years-of-european-wind-generation)
|
||||
* [ETT](https://github.com/zhouhaoyi/ETDataset)
|
||||
* [App Flow](https://github.com/alipay/Pyraformer/blob/master/data/app_zone_rpc_hour_encrypted.zip)
|
||||
|
||||
The downloaded datasets can be put in the 'data' directory. For single step forecasting, we preprocess Electricity, Wind and App Flow using scripts preprocess_elect.py, preprocess_wind.py and preprocess_flow.py respectively. You can also download preprocessed data [here](https://drive.google.com/drive/folders/1-b9tR6Tgmx48smPMetzAhVSV7-95im3X?usp=sharing). and put them in the 'data' directory. The directory structure looks like:
|
||||
|
||||
${CODE_ROOT}
|
||||
......
|
||||
|-- data
|
||||
|-- elect
|
||||
|-- test_data_elect.npy
|
||||
|-- train_data_elect.npy
|
||||
......
|
||||
|-- flow
|
||||
......
|
||||
|-- wind
|
||||
......
|
||||
|-- ETT
|
||||
|-- ETTh1.csv
|
||||
|-- ETTh2.csv
|
||||
|-- ETTm1.csv
|
||||
|-- ETTm2.csv
|
||||
|-- LD2011_2014.txt
|
||||
|-- synthetic.npy
|
||||
|
||||
Where synthetic.npy is generated by running:
|
||||
|
||||
python simulate_sin.py
|
||||
|
||||
## Training
|
||||
To perform long-range forecasting, run:
|
||||
|
||||
sh scripts/Pyraformer_LR_FC.sh
|
||||
|
||||
To perform single step forecasting, run:
|
||||
|
||||
sh scripts/Pyraformer_SS.sh
|
||||
|
||||
The meaning of each command line argument is explained in long_range_main.py and single_step_main.py, respectively.
|
||||
|
||||
## Evaluate
|
||||
Evaluation can be done by adding the -eval option to the command line. We provide pretrained models [here](https://drive.google.com/drive/folders/15av5ZhHG8tbX8HuxZNNDGBybdnuxzA83?usp=sharing). The downloaded models should be put in the 'models' directory. The directory structure is as follows:
|
||||
|
||||
${CODE_ROOT}
|
||||
......
|
||||
|-- models
|
||||
|-- LongRange
|
||||
|-- elect
|
||||
|-- 168
|
||||
|-- best_iter0.pth
|
||||
|-- best_iter1.pth
|
||||
|-- best_iter2.pth
|
||||
|-- best_iter3.pth
|
||||
|-- best_iter4.pth
|
||||
|-- 336
|
||||
......
|
||||
|-- 720
|
||||
......
|
||||
|-- ETTh1
|
||||
......
|
||||
|-- ETTm1
|
||||
......
|
||||
|-- SingleStep
|
||||
|-- elect
|
||||
|-- best_model.pth
|
||||
|-- flow
|
||||
|-- best_model.pth
|
||||
|-- wind
|
||||
|-- best_model.pth
|
||||
|
||||
Below are evaluation examples:
|
||||
|
||||
python long_range_main.py -data ETTh1 -input_size 168 -predict_step 168 -n_head 6 -eval
|
||||
|
||||
python single_step_main.py -data_path data/elect/ -dataset elect -eval
|
||||
|
||||
## Citation
|
||||
|
||||
@inproceedings{liu2022pyraformer,
|
||||
title={Pyraformer: Low-Complexity Pyramidal Attention for Long-Range Time Series Modeling and Forecasting},
|
||||
author={Liu, Shizhan and Yu, Hang and Liao, Cong and Li, Jianguo and Lin, Weiyao and Liu, Alex X and Dustdar, Schahram},
|
||||
booktitle={International Conference on Learning Representations},
|
||||
year={2022}
|
||||
}
|
||||
Vendored
-719
@@ -1,719 +0,0 @@
|
||||
import os
|
||||
import pandas as pd
|
||||
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
|
||||
from utils.tools import StandardScaler
|
||||
from utils.timefeatures import time_features
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import warnings
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
"""Long range dataloader"""
|
||||
class Dataset_ETT_hour(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None, data_path='ETTh1.csv', dataset='ETTh1', inverse=False):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24*4*4
|
||||
self.pred_len = 24*4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.pred_len = size[1]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train':0, 'val':1, 'test':2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.inverse = inverse
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12*30*24 - self.seq_len, 12*30*24+4*30*24 - self.seq_len]
|
||||
border2s = [12*30*24, 12*30*24+4*30*24, 12*30*24+8*30*24]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
data_stamp = time_features(df_stamp, timeenc=1, freq='h')
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
if self.inverse:
|
||||
self.data_y = df_data.values[border1:border2]
|
||||
else:
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end
|
||||
r_end = r_begin + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark, self.scaler.mean, self.scaler.std
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len- self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data, seq_y, mean, std):
|
||||
return self.scaler.inverse_transform(data), seq_y
|
||||
|
||||
|
||||
class Dataset_ETT_minute(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None, data_path='ETTm1.csv', dataset='ETTm1', inverse=False):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24*4*4
|
||||
self.pred_len = 24*4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.pred_len = size[1]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train':0, 'val':1, 'test':2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.inverse = inverse
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
border1s = [0, 12*30*24*4 - self.seq_len, 12*30*24*4+4*30*24*4 - self.seq_len]
|
||||
border2s = [12*30*24*4, 12*30*24*4+4*30*24*4, 12*30*24*4+8*30*24*4]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
data = self.scaler.transform(df_data.values)
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
data_stamp = time_features(df_stamp, timeenc=1, freq='h')
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
if self.inverse:
|
||||
self.data_y = df_data.values[border1:border2]
|
||||
else:
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end
|
||||
r_end = r_begin + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark, self.scaler.mean, self.scaler.std
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len- self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data, seq_y, mean, std):
|
||||
return self.scaler.inverse_transform(data), seq_y
|
||||
|
||||
|
||||
|
||||
class Dataset_Custom(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None,
|
||||
features='M', data_path='ETTh1.csv',
|
||||
target='OT', scale=True, dataset='',timeenc=0, freq='h',inverse=False):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
if size == None:
|
||||
self.seq_len = 24 * 4 * 4
|
||||
self.pred_len = 24 * 4
|
||||
else:
|
||||
self.seq_len = size[0]
|
||||
self.pred_len = size[1]
|
||||
# init
|
||||
assert flag in ['train', 'test', 'val']
|
||||
type_map = {'train': 0, 'val': 1, 'test': 2}
|
||||
self.set_type = type_map[flag]
|
||||
|
||||
self.features = features
|
||||
self.target = target
|
||||
self.scale = scale
|
||||
self.timeenc = timeenc
|
||||
self.freq = freq
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
self.__read_data__()
|
||||
|
||||
def __read_data__(self):
|
||||
self.scaler = StandardScaler()
|
||||
df_raw = pd.read_csv(os.path.join(self.root_path,
|
||||
self.data_path))
|
||||
|
||||
'''
|
||||
df_raw.columns: ['date', ...(other features), target feature]
|
||||
'''
|
||||
cols = list(df_raw.columns)
|
||||
cols.remove(self.target)
|
||||
cols.remove('date')
|
||||
df_raw = df_raw[['date'] + cols + [self.target]]
|
||||
# print(cols)
|
||||
num_train = int(len(df_raw) * 0.7)
|
||||
num_test = int(len(df_raw) * 0.2)
|
||||
num_vali = len(df_raw) - num_train - num_test
|
||||
border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len]
|
||||
border2s = [num_train, num_train + num_vali, len(df_raw)]
|
||||
border1 = border1s[self.set_type]
|
||||
border2 = border2s[self.set_type]
|
||||
|
||||
if self.features == 'M' or self.features == 'MS':
|
||||
cols_data = df_raw.columns[1:]
|
||||
df_data = df_raw[cols_data]
|
||||
elif self.features == 'S':
|
||||
df_data = df_raw[[self.target]]
|
||||
|
||||
if self.scale:
|
||||
train_data = df_data[border1s[0]:border2s[0]]
|
||||
self.scaler.fit(train_data.values)
|
||||
# print(self.scaler.mean_)
|
||||
# exit()
|
||||
data = self.scaler.transform(df_data.values)
|
||||
else:
|
||||
data = df_data.values
|
||||
|
||||
df_stamp = df_raw[['date']][border1:border2]
|
||||
df_stamp['date'] = pd.to_datetime(df_stamp.date)
|
||||
if self.timeenc == 0:
|
||||
df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
|
||||
df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
|
||||
df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
|
||||
df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
|
||||
data_stamp = df_stamp.drop(['date'], axis=1).values
|
||||
elif self.timeenc == 1:
|
||||
data_stamp = time_features(pd.to_datetime(df_stamp['date'].values), freq=self.freq)
|
||||
data_stamp = data_stamp.transpose(1, 0)
|
||||
|
||||
self.data_x = data[border1:border2]
|
||||
self.data_y = data[border1:border2]
|
||||
self.data_stamp = data_stamp
|
||||
|
||||
def __getitem__(self, index):
|
||||
s_begin = index
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end
|
||||
r_end = r_begin + self.pred_len
|
||||
|
||||
seq_x = self.data_x[s_begin:s_end]
|
||||
seq_y = self.data_y[r_begin:r_end]
|
||||
seq_x_mark = self.data_stamp[s_begin:s_end]
|
||||
seq_y_mark = self.data_stamp[r_begin:r_end]
|
||||
|
||||
return seq_x, seq_y, seq_x_mark, seq_y_mark,self.scaler.mean, self.scaler.std
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_x) - self.seq_len - self.pred_len + 1
|
||||
|
||||
def inverse_transform(self, data):
|
||||
return self.scaler.inverse_transform(data)
|
||||
|
||||
|
||||
|
||||
# """Long range dataloader for dataset elect and app flow"""
|
||||
class Dataset_Custom2(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None, data_path='ETTh1.csv', dataset='elect',
|
||||
inverse=False):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
self.seq_len = size[0]
|
||||
self.pred_len = size[1]
|
||||
# init
|
||||
assert flag in ['train', 'test']
|
||||
self.flag = flag
|
||||
|
||||
self.inverse = inverse
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
preprocess_path = os.path.join(self.root_path, self.data_path)
|
||||
self.all_data, self.covariates, self.train_end = eval('preprocess_flow')(preprocess_path)
|
||||
self.all_data = torch.from_numpy(self.all_data).transpose(0, 1)
|
||||
self.covariates = torch.from_numpy(self.covariates)
|
||||
self.test_start = self.train_end - self.seq_len + 1
|
||||
self.window_stride = 24
|
||||
self.seq_num = self.all_data.size(0)
|
||||
|
||||
def fit(self, data):
|
||||
mean = data.mean()
|
||||
std = data.std()
|
||||
return mean, std
|
||||
|
||||
def inverse_transform(self, output, seq_y, mean, std):
|
||||
output = output * (mean.unsqueeze(1).unsqueeze(1) + 1)
|
||||
seq_y = seq_y * (mean.unsqueeze(1).unsqueeze(1) + 1)
|
||||
return output, seq_y
|
||||
|
||||
def __len__(self):
|
||||
if self.flag == 'train':
|
||||
self.window_per_seq = (self.train_end - self.seq_len - self.pred_len) // self.window_stride
|
||||
return self.window_per_seq * self.seq_num
|
||||
else:
|
||||
self.window_per_seq = (self.all_data.size(1) - self.test_start - self.seq_len - self.pred_len) // self.window_stride
|
||||
return self.window_per_seq * self.seq_num
|
||||
|
||||
def __getitem__(self, index):
|
||||
seq_idx = index // self.window_per_seq
|
||||
window_idx = index % self.window_per_seq
|
||||
|
||||
if self.flag == 'train':
|
||||
s_begin = window_idx * self.window_stride
|
||||
else:
|
||||
s_begin = self.test_start + window_idx * self.window_stride
|
||||
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end
|
||||
r_end = r_begin + self.pred_len
|
||||
|
||||
seq_x = self.all_data[seq_idx, s_begin:s_end].clone()
|
||||
seq_y = self.all_data[seq_idx, r_begin:r_end].clone()
|
||||
mean, std = self.fit(seq_x)
|
||||
if mean > 0:
|
||||
seq_x = seq_x / (mean + 1)
|
||||
seq_y = seq_y / (mean + 1)
|
||||
|
||||
if len(self.covariates.size()) == 2:
|
||||
seq_x_mark = self.covariates[s_begin:s_end]
|
||||
seq_x_mark[:, -1] = int(seq_idx)
|
||||
seq_y_mark = self.covariates[r_begin:r_end]
|
||||
seq_y_mark[:, -1] = int(seq_idx)
|
||||
else:
|
||||
seq_x_mark = self.covariates[s_begin:s_end, seq_idx]
|
||||
seq_x_mark[:, -1] = int(seq_idx)
|
||||
seq_y_mark = self.covariates[r_begin:r_end, seq_idx]
|
||||
seq_y_mark[:, -1] = int(seq_idx)
|
||||
|
||||
return seq_x.unsqueeze(1), seq_y.unsqueeze(1), seq_x_mark, seq_y_mark, mean, std
|
||||
|
||||
|
||||
"""Long range dataloader for synthetic dataset"""
|
||||
class Dataset_Synthetic(Dataset):
|
||||
def __init__(self, root_path, flag='train', size=None, data_path='synthetic.npy', dataset='synthetic', inverse=False):
|
||||
# size [seq_len, label_len, pred_len]
|
||||
# info
|
||||
self.seq_len = size[0]
|
||||
self.pred_len = size[1]
|
||||
# init
|
||||
assert flag in ['train', 'test']
|
||||
self.flag = flag
|
||||
self.inverse = inverse
|
||||
|
||||
self.root_path = root_path
|
||||
self.data_path = data_path
|
||||
preprocess_path = os.path.join(self.root_path, self.data_path)
|
||||
self.all_data =np.load(preprocess_path)
|
||||
self.all_data = torch.from_numpy(self.all_data)
|
||||
self.all_data, self.covariates = self.all_data[:, :, 0], self.all_data[:, :, 1:]
|
||||
self.seq_num = self.all_data.size(0)
|
||||
|
||||
self.window_stride = 24
|
||||
window_per_seq = (self.all_data.shape[1] - self.seq_len - self.pred_len) / self.window_stride
|
||||
self.train_end = self.seq_len + self.pred_len + int(0.9 * window_per_seq) * self.window_stride
|
||||
self.test_start = self.train_end - self.seq_len + 1
|
||||
|
||||
def fit(self, data):
|
||||
mean = data.mean()
|
||||
std = data.std()
|
||||
return mean, std
|
||||
|
||||
def inverse_transform(self, output, seq_y, mean, std):
|
||||
output = output * (mean.unsqueeze(1).unsqueeze(1) + 1)
|
||||
seq_y = seq_y * (mean.unsqueeze(1).unsqueeze(1) + 1)
|
||||
return output, seq_y
|
||||
|
||||
def __len__(self):
|
||||
if self.flag == 'train':
|
||||
self.window_per_seq = (self.train_end - self.seq_len - self.pred_len) // self.window_stride
|
||||
return self.window_per_seq * self.seq_num
|
||||
else:
|
||||
self.window_per_seq = (self.all_data.size(1) - self.test_start - self.seq_len - self.pred_len) // self.window_stride
|
||||
return self.window_per_seq * self.seq_num
|
||||
|
||||
def __getitem__(self, index):
|
||||
seq_idx = index // self.window_per_seq
|
||||
window_idx = index % self.window_per_seq
|
||||
|
||||
if self.flag == 'train':
|
||||
s_begin = window_idx * self.window_stride
|
||||
else:
|
||||
s_begin = self.test_start + window_idx * self.window_stride
|
||||
|
||||
s_end = s_begin + self.seq_len
|
||||
r_begin = s_end
|
||||
r_end = r_begin + self.pred_len
|
||||
|
||||
seq_x = self.all_data[seq_idx, s_begin:s_end].clone()
|
||||
seq_y = self.all_data[seq_idx, r_begin:r_end].clone()
|
||||
|
||||
mean, std = self.fit(seq_x)
|
||||
if mean > 0:
|
||||
seq_x = seq_x / (mean + 1)
|
||||
seq_y = seq_y / (mean + 1)
|
||||
|
||||
seq_x_mark = self.covariates[seq_idx, s_begin:s_end]
|
||||
seq_y_mark = self.covariates[seq_idx, r_begin:r_end]
|
||||
|
||||
return seq_x.unsqueeze(1), seq_y.unsqueeze(1), seq_x_mark, seq_y_mark, mean, std
|
||||
|
||||
|
||||
def get_all_v(train_data, train_end, seq_len, pred_len, window_stride, type):
|
||||
"""Get the normalization parameters of each sequence"""
|
||||
seq_num = train_data.size(0)
|
||||
window_per_seq = (train_end - seq_len - pred_len) // window_stride
|
||||
window_number = seq_num * window_per_seq
|
||||
|
||||
v = torch.zeros(window_number, dtype=torch.float64)
|
||||
for index in range(window_number):
|
||||
seq_idx = index // window_per_seq
|
||||
window_idx = index % window_per_seq
|
||||
|
||||
s_begin = window_idx * window_stride
|
||||
s_end = s_begin + seq_len
|
||||
|
||||
seq_x = train_data[seq_idx, s_begin:s_end].clone()
|
||||
if type == 'mean':
|
||||
mean = seq_x.mean()
|
||||
v[index] = mean + 1
|
||||
else:
|
||||
std = seq_x.std()
|
||||
v[index] = std
|
||||
|
||||
return v
|
||||
|
||||
|
||||
def gen_covariates(times, num_covariates):
|
||||
"""Get covariates"""
|
||||
covariates = np.zeros((times.shape[0], num_covariates))
|
||||
for i, input_time in enumerate(times):
|
||||
covariates[i, 0] = input_time.weekday() / 7
|
||||
covariates[i, 1] = input_time.hour / 24
|
||||
covariates[i, 2] = input_time.month / 12
|
||||
|
||||
return covariates
|
||||
|
||||
|
||||
def preprocess_elect(csv_path):
|
||||
"""preprocess the elect dataset for long range forecasting"""
|
||||
num_covariates = 4
|
||||
train_start = '2011-01-01 00:00:00'
|
||||
train_end = '2014-04-01 23:00:00'
|
||||
test_start = '2014-04-01 00:00:00'
|
||||
test_end = '2014-09-07 23:00:00'
|
||||
|
||||
data_frame = pd.read_csv(csv_path, sep=";", index_col=0, parse_dates=True, decimal=',')
|
||||
data_frame = data_frame.resample('1H',label = 'left',closed = 'right').sum()[train_start:test_end]
|
||||
data_frame.fillna(0, inplace=True)
|
||||
|
||||
covariates = gen_covariates(data_frame[train_start:test_end].index, num_covariates)
|
||||
all_data = data_frame[train_start:test_end].values
|
||||
data_start = (all_data!=0).argmax(axis=0) #find first nonzero value in each time series
|
||||
train_end = len(data_frame[train_start:train_end].values)
|
||||
|
||||
all_data = all_data[:, data_start < 10000]
|
||||
data_start = data_start[data_start < 10000]
|
||||
split_start = data_start.max()
|
||||
all_data = all_data[split_start:]
|
||||
covariates = covariates[split_start:]
|
||||
train_end = train_end - split_start
|
||||
|
||||
return all_data.astype(np.float32), covariates.astype(np.float32), train_end
|
||||
|
||||
|
||||
def preprocess_flow(csv_path):
|
||||
"""preprocess the app flow dataset for long range forecasting"""
|
||||
data_frame = pd.read_csv(csv_path, names=['app_name', 'zone', 'time', 'value'], parse_dates=True)
|
||||
grouped_data = list(data_frame.groupby(["app_name", "zone"]))
|
||||
# covariates = gen_covariates(data_frame.index, 3)
|
||||
all_data = []
|
||||
min_length = 10000
|
||||
for i in range(len(grouped_data)):
|
||||
single_df = grouped_data[i][1].drop(labels=['app_name', 'zone'], axis=1).sort_values(by="time", ascending=True)
|
||||
times = pd.to_datetime(single_df.time)
|
||||
single_df['weekday'] = times.dt.dayofweek / 7
|
||||
single_df['hour'] = times.dt.hour / 24
|
||||
single_df['month'] = times.dt.month / 12
|
||||
temp_data = single_df.values[:, 1:]
|
||||
if (temp_data[:, 0] == 0).sum() / len(temp_data) > 0.2 or len(temp_data) < 3000:
|
||||
continue
|
||||
|
||||
if len(temp_data) < min_length:
|
||||
min_length = len(temp_data)
|
||||
|
||||
all_data.append(temp_data)
|
||||
|
||||
all_data = np.array([data[len(data)-min_length:, :] for data in all_data]).transpose(1, 0, 2).astype(np.float32)
|
||||
train_end = min(int(0.8 * min_length), min_length - 1000)
|
||||
covariates = all_data.copy()
|
||||
covariates[:, :, :-1] = covariates[:, :, 1:]
|
||||
|
||||
return all_data[:, :, 0], covariates, train_end
|
||||
|
||||
|
||||
"""Single step dataloader"""
|
||||
def split(split_start, label, cov, pred_length):
|
||||
all_data = []
|
||||
for batch_idx in range(len(label)):
|
||||
batch_label = label[batch_idx]
|
||||
for i in range(pred_length):
|
||||
single_data = batch_label[i:(split_start+i)].clone().unsqueeze(1)
|
||||
single_data[-1] = -1
|
||||
single_cov = cov[batch_idx, i:(split_start+i), :].clone()
|
||||
temp_data = [single_data, single_cov]
|
||||
single_data = torch.cat(temp_data, dim=1)
|
||||
all_data.append(single_data)
|
||||
data = torch.stack(all_data, dim=0)
|
||||
label = label[:, -pred_length:].reshape(pred_length*len(label))
|
||||
|
||||
return data, label
|
||||
|
||||
|
||||
"""Single step training dataloader for the electricity dataset"""
|
||||
class electTrainDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length, batch_size):
|
||||
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
|
||||
|
||||
# Resample windows according to the average amplitude
|
||||
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
|
||||
weights = torch.as_tensor(np.abs(v[:,0])/np.sum(np.abs(v[:,0])), dtype=torch.double)
|
||||
num_samples = weights.size(0)
|
||||
sample_index = torch.multinomial(weights, num_samples, True)
|
||||
self.data = self.data[sample_index]
|
||||
|
||||
self.label = torch.from_numpy(np.load(os.path.join(data_path, f'train_label_{data_name}.npy')))
|
||||
self.label = self.label[sample_index]
|
||||
|
||||
self.train_len = len(self.data) // batch_size
|
||||
self.pred_length = predict_length
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __len__(self):
|
||||
return self.train_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
if (index+1) <= self.train_len:
|
||||
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
|
||||
label = self.label[index*self.batch_size:(index+1)*self.batch_size].clone()
|
||||
else:
|
||||
all_data = self.data[index*self.batch_size:].clone()
|
||||
label = self.label[index*self.batch_size:].clone()
|
||||
|
||||
cov = all_data[:, :, 2:]
|
||||
|
||||
split_start = len(label[0]) - self.pred_length + 1
|
||||
data, label = split(split_start, label, cov, self.pred_length)
|
||||
|
||||
return data, label
|
||||
|
||||
|
||||
"""Single step testing dataloader for the electricity dataset"""
|
||||
class electTestDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length):
|
||||
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
|
||||
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
|
||||
self.label = np.load(os.path.join(data_path, f'test_label_{data_name}.npy'))
|
||||
self.test_len = self.data.shape[0]
|
||||
self.pred_length = predict_length
|
||||
|
||||
def __len__(self):
|
||||
return self.test_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
all_data = torch.from_numpy(self.data[index].copy())
|
||||
cov = all_data[:, 2:]
|
||||
label = torch.from_numpy(self.label[index].copy())
|
||||
v = float(self.v[index][0])
|
||||
if v > 0:
|
||||
data = label / v
|
||||
else:
|
||||
data = label
|
||||
|
||||
split_start = len(label) - self.pred_length + 1
|
||||
all_data = []
|
||||
for i in range(self.pred_length):
|
||||
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
|
||||
single_data[-1] = -1
|
||||
single_cov = cov[i:(split_start+i), :].clone()
|
||||
single_data = torch.cat([single_data, single_cov], dim=1)
|
||||
all_data.append(single_data)
|
||||
all_data = torch.stack(all_data, dim=0)
|
||||
label = label[-self.pred_length:]
|
||||
|
||||
return all_data, label, v
|
||||
|
||||
|
||||
"""Single step training dataloader for the app flow dataset"""
|
||||
class flowTrainDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length, batch_size):
|
||||
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
|
||||
|
||||
# Resample windows according to the average amplitude
|
||||
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
|
||||
weights = torch.as_tensor(np.abs(v)/np.sum(np.abs(v)), dtype=torch.double)
|
||||
num_samples = weights.size(0)
|
||||
sample_index = torch.multinomial(weights, num_samples, True)
|
||||
self.data = self.data[sample_index]
|
||||
|
||||
self.label = self.data[:, :, 0]
|
||||
|
||||
self.train_len = len(self.data) // batch_size
|
||||
self.pred_length = predict_length
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __len__(self):
|
||||
return self.train_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
if (index+1) <= self.train_len:
|
||||
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
|
||||
label = self.label[index*self.batch_size:(index+1)*self.batch_size].clone()
|
||||
else:
|
||||
all_data = self.data[index*self.batch_size:].clone()
|
||||
label = self.label[index*self.batch_size:].clone()
|
||||
|
||||
cov = all_data[:, :, 1:]
|
||||
|
||||
split_start = len(label[0]) - self.pred_length + 1
|
||||
data, label = split(split_start, label, cov, self.pred_length)
|
||||
|
||||
return data, label
|
||||
|
||||
|
||||
"""Single step testing dataloader for the all flow dataset"""
|
||||
class flowTestDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length):
|
||||
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
|
||||
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
|
||||
self.label = self.data
|
||||
self.test_len = self.data.shape[0]
|
||||
self.pred_length = predict_length
|
||||
|
||||
def __len__(self):
|
||||
return self.test_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
all_data = torch.from_numpy(self.data[index].copy())
|
||||
cov = all_data[:, 1:]
|
||||
data = all_data[:, 0]
|
||||
label = torch.from_numpy(self.label[index, :, 0].copy())
|
||||
v = float(self.v[index])
|
||||
|
||||
split_start = len(label) - self.pred_length + 1
|
||||
all_data = []
|
||||
for i in range(self.pred_length):
|
||||
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
|
||||
single_data[-1] = -1
|
||||
single_cov = cov[i:(split_start+i), :].clone()
|
||||
single_data = torch.cat([single_data, single_cov], dim=1)
|
||||
all_data.append(single_data)
|
||||
all_data = torch.stack(all_data, dim=0)
|
||||
label = label[-self.pred_length:] * v
|
||||
|
||||
return all_data, label, v
|
||||
|
||||
|
||||
"""Single step training dataloader for the wind dataset"""
|
||||
class windTrainDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length, batch_size):
|
||||
self.data = torch.from_numpy(np.load(os.path.join(data_path, f'train_data_{data_name}.npy')))
|
||||
|
||||
# Resample windows according to the average amplitude
|
||||
v = np.load(os.path.join(data_path, f'train_v_{data_name}.npy'))
|
||||
weights = torch.as_tensor(np.abs(v)/np.sum(np.abs(v)), dtype=torch.double)
|
||||
num_samples = weights.size(0)
|
||||
sample_index = torch.multinomial(weights, num_samples, True)
|
||||
self.data = self.data[sample_index]
|
||||
|
||||
self.train_len = len(self.data) // batch_size
|
||||
self.pred_length = predict_length
|
||||
self.batch_size = batch_size
|
||||
|
||||
def __len__(self):
|
||||
return self.train_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
if (index+1) <= self.train_len:
|
||||
all_data = self.data[index*self.batch_size:(index+1)*self.batch_size].clone()
|
||||
else:
|
||||
all_data = self.data[index*self.batch_size:].clone()
|
||||
|
||||
cov = all_data[:, :, 1:]
|
||||
label = all_data[:, :, 0]
|
||||
|
||||
split_start = len(label[0]) - self.pred_length + 1
|
||||
data, label = split(split_start, label, cov, self.pred_length)
|
||||
|
||||
return data, label
|
||||
|
||||
|
||||
"""Single step testing dataloader for the wind dataset"""
|
||||
class windTestDataset(Dataset):
|
||||
def __init__(self, data_path, data_name, predict_length):
|
||||
self.data = np.load(os.path.join(data_path, f'test_data_{data_name}.npy'))
|
||||
self.v = np.load(os.path.join(data_path, f'test_v_{data_name}.npy'))
|
||||
self.test_len = self.data.shape[0]
|
||||
self.pred_length = predict_length
|
||||
|
||||
def __len__(self):
|
||||
return self.test_len
|
||||
|
||||
def __getitem__(self, index):
|
||||
all_data = torch.from_numpy(self.data[index].copy())
|
||||
cov = all_data[:, 1:]
|
||||
data = all_data[:, 0]
|
||||
v = float(self.v[index])
|
||||
label = data * v
|
||||
|
||||
split_start = len(label) - self.pred_length + 1
|
||||
all_data = []
|
||||
for i in range(self.pred_length):
|
||||
single_data = data[i:(split_start+i)].clone().unsqueeze(1)
|
||||
single_data[-1] = -1
|
||||
single_cov = cov[i:(split_start+i), :].clone()
|
||||
single_data = torch.cat([single_data, single_cov], dim=1)
|
||||
all_data.append(single_data)
|
||||
all_data = torch.stack(all_data, dim=0)
|
||||
label = label[-self.pred_length:]
|
||||
|
||||
return all_data, label, v
|
||||
|
||||
Vendored
-409
@@ -1,409 +0,0 @@
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import time
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
|
||||
import pyraformer.Pyraformer_LR as Pyraformer
|
||||
from tqdm import tqdm
|
||||
from data_loader import *
|
||||
from utils.tools import TopkMSELoss, metric
|
||||
|
||||
|
||||
def prepare_dataloader(args):
|
||||
""" Load data and prepare dataloader. """
|
||||
|
||||
data_dict = {
|
||||
'ETTh1':Dataset_ETT_hour,
|
||||
'ETTh2':Dataset_ETT_hour,
|
||||
'ETTm1':Dataset_ETT_minute,
|
||||
'ETTm2':Dataset_ETT_minute,
|
||||
'electricity':Dataset_Custom,
|
||||
'exchange':Dataset_Custom,
|
||||
'traffic':Dataset_Custom,
|
||||
'weather':Dataset_Custom,
|
||||
'ili':Dataset_Custom,
|
||||
# 'flow': Dataset_Custom2,
|
||||
# 'synthetic': Dataset_Synthetic,
|
||||
}
|
||||
Data = data_dict[args.data]
|
||||
|
||||
# prepare training dataset and dataloader
|
||||
shuffle_flag = True; drop_last = True; batch_size = args.batch_size
|
||||
train_set = Data(
|
||||
root_path=args.root_path,
|
||||
data_path=args.data_path,
|
||||
flag='train',
|
||||
size=[args.input_size, args.predict_step],
|
||||
inverse=args.inverse,
|
||||
dataset=args.data
|
||||
)
|
||||
print('train', len(train_set))
|
||||
train_loader = DataLoader(
|
||||
train_set,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle_flag,
|
||||
num_workers=0,
|
||||
drop_last=drop_last)
|
||||
|
||||
# prepare testing dataset and dataloader
|
||||
shuffle_flag = False; drop_last = False; batch_size = args.batch_size
|
||||
test_set = Data(
|
||||
root_path=args.root_path,
|
||||
data_path=args.data_path,
|
||||
flag='test',
|
||||
size=[args.input_size, args.predict_step],
|
||||
inverse=args.inverse,
|
||||
dataset=args.data
|
||||
)
|
||||
print('test', len(test_set))
|
||||
test_loader = DataLoader(
|
||||
test_set,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle_flag,
|
||||
num_workers=0,
|
||||
drop_last=drop_last)
|
||||
|
||||
return train_loader, train_set, test_loader, test_set
|
||||
|
||||
|
||||
def sample_mining_scheduler(epoch, batch_size):
|
||||
if epoch < 2:
|
||||
topk = batch_size
|
||||
elif epoch < 4:
|
||||
topk = int(batch_size * (5 - epoch) / (6 - epoch))
|
||||
else:
|
||||
topk = int(0.5 * batch_size)
|
||||
|
||||
return topk
|
||||
|
||||
|
||||
def dataset_parameters(args, dataset):
|
||||
"""Prepare specific parameters for different datasets"""
|
||||
dataset2enc_in = {
|
||||
'ETTh1':7,
|
||||
'ETTh2':7,
|
||||
'ETTm1':7,
|
||||
'ETTm2':7,
|
||||
'electricity':321,
|
||||
'exchange':8,
|
||||
'traffic':862,
|
||||
'weather':21,
|
||||
'ili':7,
|
||||
'flow': 1,
|
||||
'synthetic': 1
|
||||
}
|
||||
dataset2cov_size = {
|
||||
'ETTh1':4,
|
||||
'ETTh2':4,
|
||||
'ETTm1':4,
|
||||
'ETTm2':4,
|
||||
'electricity':4,
|
||||
'exchange':4,
|
||||
'traffic':4,
|
||||
'weather':4,
|
||||
'ili':4,
|
||||
'elect':3,
|
||||
'flow': 3,
|
||||
'synthetic': 3,
|
||||
}
|
||||
dataset2seq_num = {
|
||||
'ETTh1':1,
|
||||
'ETTh2':1,
|
||||
'ETTm1':1,
|
||||
'ETTm2':1,
|
||||
'electricity':1,
|
||||
'exchange':1,
|
||||
'traffic':1,
|
||||
'weather':1,
|
||||
'ili':1,
|
||||
'elect':321,
|
||||
'flow': 1077,
|
||||
'synthetic': 60
|
||||
}
|
||||
dataset2embed = {
|
||||
'ETTh1':'DataEmbedding',
|
||||
'ETTh2':'DataEmbedding',
|
||||
'ETTm1':'DataEmbedding',
|
||||
'ETTm2':'DataEmbedding',
|
||||
'elect':'CustomEmbedding',
|
||||
'electricity':'CustomEmbedding',
|
||||
'exchange':'CustomEmbedding',
|
||||
'traffic':'CustomEmbedding',
|
||||
'weather':'CustomEmbedding',
|
||||
'ili':'CustomEmbedding',
|
||||
'flow': 'CustomEmbedding',
|
||||
'synthetic': 'CustomEmbedding'
|
||||
}
|
||||
|
||||
args.enc_in = dataset2enc_in[dataset]
|
||||
args.dec_in = dataset2enc_in[dataset]
|
||||
args.covariate_size = dataset2cov_size[dataset]
|
||||
args.seq_num = dataset2seq_num[dataset]
|
||||
args.embed_type = dataset2embed[dataset]
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def train_epoch(model, train_dataset, training_loader, optimizer, opt, epoch):
|
||||
""" Epoch operation in training phase. """
|
||||
|
||||
model.train()
|
||||
total_loss = 0
|
||||
total_pred_number = 0
|
||||
warm = False
|
||||
for batch in tqdm(training_loader, mininterval=2,
|
||||
desc=' - (Training) ', leave=False):
|
||||
# prepare data
|
||||
batch_x, batch_y, batch_x_mark, batch_y_mark, mean, std = map(lambda x: x.float().to(opt.device), batch)
|
||||
dec_inp = torch.zeros_like(batch_y).float()
|
||||
optimizer.zero_grad()
|
||||
|
||||
# forward
|
||||
if opt.decoder == 'attention':
|
||||
if opt.pretrain and epoch < 1:
|
||||
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, True)
|
||||
batch_y = torch.cat([batch_x, batch_y], dim=1)
|
||||
else:
|
||||
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
|
||||
elif opt.decoder == 'FC':
|
||||
# Add a predict token into the history sequence
|
||||
predict_token = torch.zeros(batch_x.size(0), 1, batch_x.size(-1), device=batch_x.device)
|
||||
batch_x = torch.cat([batch_x, predict_token], dim=1)
|
||||
batch_x_mark = torch.cat([batch_x_mark, batch_y_mark[:, 0:1, :]], dim=1)
|
||||
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
|
||||
|
||||
# determine the loss function
|
||||
if opt.hard_sample_mining and not (opt.pretrain and epoch < 1):
|
||||
topk = sample_mining_scheduler(epoch, batch_x.size(0))
|
||||
criterion = TopkMSELoss(topk)
|
||||
else:
|
||||
criterion = torch.nn.MSELoss(reduction='none')
|
||||
# if inverse, both the output and the ground truth are denormalized.
|
||||
if opt.inverse:
|
||||
outputs, batch_y = train_dataset.inverse_transform(outputs, batch_y, mean, std)
|
||||
# compute loss
|
||||
losses = criterion(outputs, batch_y)
|
||||
loss = losses.mean()
|
||||
loss.backward()
|
||||
|
||||
""" update parameters """
|
||||
optimizer.step()
|
||||
total_loss += losses.sum().item()
|
||||
total_pred_number += losses.numel()
|
||||
|
||||
return total_loss / total_pred_number
|
||||
|
||||
|
||||
def eval_epoch(model, test_dataset, test_loader, opt, epoch):
|
||||
""" Epoch operation in evaluation phase. """
|
||||
|
||||
model.eval()
|
||||
preds = []
|
||||
trues = []
|
||||
warm = 0
|
||||
with torch.no_grad():
|
||||
for batch in tqdm(test_loader, mininterval=2,
|
||||
desc=' - (Validation) ', leave=False):
|
||||
""" prepare data """
|
||||
batch_x, batch_y, batch_x_mark, batch_y_mark, mean, std = map(lambda x: x.float().to(opt.device), batch)
|
||||
dec_inp = torch.zeros_like(batch_y).float()
|
||||
|
||||
# forward
|
||||
if opt.decoder == 'FC':
|
||||
# Add a predict token into the history sequence
|
||||
predict_token = torch.zeros(batch_x.size(0), 1, batch_x.size(-1), device=batch_x.device)
|
||||
batch_x = torch.cat([batch_x, predict_token], dim=1)
|
||||
batch_x_mark = torch.cat([batch_x_mark, batch_y_mark[:, 0:1, :]], dim=1)
|
||||
outputs = model(batch_x, batch_x_mark, dec_inp, batch_y_mark, False)
|
||||
|
||||
warm += 1
|
||||
# if inverse, both the output and the ground truth are denormalized.
|
||||
if opt.inverse:
|
||||
outputs, batch_y = test_dataset.inverse_transform(outputs, batch_y, mean, std)
|
||||
|
||||
pred = outputs.detach().cpu().numpy()
|
||||
true = batch_y.detach().cpu().numpy()
|
||||
|
||||
preds.append(pred)
|
||||
trues.append(true)
|
||||
|
||||
preds = np.array(preds)
|
||||
|
||||
trues = np.array(trues)
|
||||
# preds = preds.reshape(-1, preds.shape[-2], preds.shape[-1])
|
||||
# trues = trues.reshape(-1, trues.shape[-2], trues.shape[-1])
|
||||
preds = np.concatenate(preds, axis=0)
|
||||
print(preds.shape)
|
||||
trues = np.concatenate(trues, axis=0)
|
||||
# np.save('./results/' + 'pred.npy', preds)
|
||||
# np.save('./results/'+ 'true.npy', trues)
|
||||
print('test shape:{}'.format(preds.shape))
|
||||
mae, mse, rmse, mape, mspe = metric(preds, trues)
|
||||
print('Epoch {}, mse:{}, mae:{}, rmse:{}, mape:{}, mspe:{}'.format(epoch, mse, mae, rmse, mape, mspe))
|
||||
|
||||
return mse, mae, rmse, mape, mspe
|
||||
|
||||
|
||||
def train(model, optimizer, scheduler, opt, model_save_dir):
|
||||
""" Start training. """
|
||||
|
||||
best_mse = 100000000
|
||||
|
||||
""" prepare dataloader """
|
||||
training_dataloader, train_dataset, test_dataloader, test_dataset = prepare_dataloader(opt)
|
||||
|
||||
|
||||
best_metrics = []
|
||||
for epoch_i in range(opt.epoch):
|
||||
epoch = epoch_i + 1
|
||||
print('[ Epoch', epoch, ']')
|
||||
|
||||
start = time.time()
|
||||
train_mse = train_epoch(model, train_dataset, training_dataloader, optimizer, opt, epoch_i)
|
||||
print(' - (Training) '
|
||||
'MSE: {mse: 8.5f}'
|
||||
'elapse: {elapse:3.3f} min'
|
||||
.format(mse=train_mse, elapse=(time.time() - start) / 60))
|
||||
|
||||
mse, mae, rmse, mape, mspe = eval_epoch(model, test_dataset, test_dataloader, opt, epoch_i)
|
||||
|
||||
scheduler.step()
|
||||
|
||||
current_metrics = [float(mse), float(mae), float(rmse), float(mape), float(mspe)]
|
||||
if best_mse > mse:
|
||||
best_mse = mse
|
||||
best_metrics = current_metrics
|
||||
torch.save(
|
||||
{
|
||||
"state_dict": model.state_dict(),
|
||||
"metrics": best_metrics
|
||||
},
|
||||
model_save_dir
|
||||
)
|
||||
|
||||
return best_metrics
|
||||
|
||||
|
||||
def evaluate(model, opt, model_save_dir):
|
||||
"""Evaluate preptrained models"""
|
||||
best_mse = 100000000
|
||||
|
||||
""" prepare dataloader """
|
||||
_, _, test_dataloader, test_dataset = prepare_dataloader(opt)
|
||||
""" load pretrained model """
|
||||
checkpoint = torch.load(model_save_dir)["state_dict"]
|
||||
model.load_state_dict(checkpoint)
|
||||
|
||||
best_metrics = []
|
||||
mse, mae, rmse, mape, mspe = eval_epoch(model, test_dataset, test_dataloader, opt, 0)
|
||||
|
||||
current_metrics = [float(mse), float(mae), float(rmse), float(mape), float(mspe)]
|
||||
if best_mse > mse:
|
||||
best_mse = mse
|
||||
best_metrics = current_metrics
|
||||
|
||||
return best_metrics
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
# running mode
|
||||
parser.add_argument('-eval', action='store_true', default=False)
|
||||
|
||||
# Path parameters
|
||||
parser.add_argument('-data', type=str, default='ETTh1')
|
||||
parser.add_argument('-root_path', type=str, default='../dataset/', help='root path of the data file')
|
||||
parser.add_argument('-data_path', type=str, default='ETTh1.csv', help='data file')
|
||||
|
||||
# Dataloader parameters.
|
||||
parser.add_argument('-input_size', type=int, default=168)
|
||||
parser.add_argument('-predict_step', type=int, default=168)
|
||||
parser.add_argument('-inverse', action='store_true', help='denormalize output data', default=False)
|
||||
|
||||
# Architecture selection.
|
||||
parser.add_argument('-model', type=str, default='Pyraformer')
|
||||
parser.add_argument('-decoder', type=str, default='FC') # selection: [FC, attention]
|
||||
|
||||
# Training parameters.
|
||||
parser.add_argument('-epoch', type=int, default=5)
|
||||
parser.add_argument('-batch_size', type=int, default=32)
|
||||
parser.add_argument('-pretrain', action='store_true', default=False)
|
||||
parser.add_argument('-hard_sample_mining', action='store_true', default=False)
|
||||
parser.add_argument('-dropout', type=float, default=0.05)
|
||||
parser.add_argument('-lr', type=float, default=1e-4)
|
||||
parser.add_argument('-lr_step', type=float, default=0.1)
|
||||
|
||||
# Common Model parameters.
|
||||
parser.add_argument('-d_model', type=int, default=512)
|
||||
parser.add_argument('-d_inner_hid', type=int, default=512)
|
||||
parser.add_argument('-d_k', type=int, default=128)
|
||||
parser.add_argument('-d_v', type=int, default=128)
|
||||
parser.add_argument('-d_bottleneck', type=int, default=128)
|
||||
parser.add_argument('-n_head', type=int, default=4)
|
||||
parser.add_argument('-n_layer', type=int, default=4)
|
||||
|
||||
# Pyraformer parameters.
|
||||
parser.add_argument('-window_size', type=str, default='[4, 4, 4]') # The number of children of a parent node.
|
||||
parser.add_argument('-inner_size', type=int, default=3) # The number of ajacent nodes.
|
||||
# CSCM structure. selection: [Bottleneck_Construct, Conv_Construct, MaxPooling_Construct, AvgPooling_Construct]
|
||||
parser.add_argument('-CSCM', type=str, default='Bottleneck_Construct')
|
||||
parser.add_argument('-truncate', action='store_true', default=False) # Whether to remove coarse-scale nodes from the attention structure
|
||||
parser.add_argument('-use_tvm', action='store_true', default=False) # Whether to use TVM.
|
||||
|
||||
# Experiment repeat times.
|
||||
parser.add_argument('-iter_num', type=int, default=1) # Repeat number.
|
||||
|
||||
opt = parser.parse_args()
|
||||
return opt
|
||||
|
||||
|
||||
def main(opt, iter_index):
|
||||
""" Main function. """
|
||||
print('[Info] parameters: {}'.format(opt))
|
||||
|
||||
if torch.cuda.is_available():
|
||||
opt.device = torch.device("cuda")
|
||||
else:
|
||||
opt.device = torch.device('cpu')
|
||||
|
||||
""" prepare model """
|
||||
model = eval(opt.model).Model(opt)
|
||||
|
||||
model.to(opt.device)
|
||||
|
||||
""" number of parameters """
|
||||
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
print('[Info] Number of parameters: {}'.format(num_params))
|
||||
|
||||
""" train or evaluate the model """
|
||||
model_save_dir = 'models/LongRange/{}/{}/'.format(opt.data, opt.predict_step)
|
||||
os.makedirs(model_save_dir, exist_ok=True)
|
||||
model_save_dir += 'best_iter{}.pth'.format(iter_index)
|
||||
if opt.eval:
|
||||
best_metrics = evaluate(model, opt, model_save_dir)
|
||||
else:
|
||||
""" optimizer and scheduler """
|
||||
optimizer = optim.Adam(filter(lambda x: x.requires_grad, model.parameters()), opt.lr)
|
||||
scheduler = optim.lr_scheduler.StepLR(optimizer, 1, gamma=opt.lr_step)
|
||||
best_metrics = train(model, optimizer, scheduler, opt, model_save_dir)
|
||||
|
||||
print('Iteration best metrics: {}'.format(best_metrics))
|
||||
return best_metrics
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
opt = parse_args()
|
||||
opt = dataset_parameters(opt, opt.data)
|
||||
opt.window_size = eval(opt.window_size)
|
||||
iter_num = opt.iter_num
|
||||
all_perf = []
|
||||
for i in range(iter_num):
|
||||
metrics = main(opt, i)
|
||||
all_perf.append(metrics)
|
||||
all_perf = np.array(all_perf)
|
||||
all_perf = all_perf.mean(0)
|
||||
print('Average Metrics: {}'.format(all_perf))
|
||||
|
||||
Vendored
-121
@@ -1,121 +0,0 @@
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
import pandas as pd
|
||||
import math
|
||||
import numpy as np
|
||||
import random
|
||||
from tqdm import trange
|
||||
|
||||
from io import BytesIO
|
||||
from urllib.request import urlopen
|
||||
from zipfile import ZipFile
|
||||
|
||||
from math import sqrt
|
||||
from pandas import read_csv, DataFrame
|
||||
from scipy import stats
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use('Agg')
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
def prep_data(data, covariates, data_start, train = True):
|
||||
"""Divide the training sequence into windows"""
|
||||
time_len = data.shape[0]
|
||||
input_size = window_size-stride_size
|
||||
windows_per_series = np.full((num_series), (time_len-input_size) // stride_size)
|
||||
if train: windows_per_series -= (data_start+stride_size-1) // stride_size
|
||||
total_windows = np.sum(windows_per_series)
|
||||
x_input = np.zeros((total_windows, window_size, 1 + num_covariates + 1), dtype='float32')
|
||||
label = np.zeros((total_windows, window_size), dtype='float32')
|
||||
v_input = np.zeros((total_windows, 2), dtype='float32')
|
||||
count = 0
|
||||
if not train:
|
||||
covariates = covariates[-time_len:]
|
||||
for series in trange(num_series):
|
||||
cov_age = stats.zscore(np.arange(total_time-data_start[series])) # shape:(series_len,)
|
||||
if train:
|
||||
covariates[data_start[series]:time_len, 0] = cov_age[:time_len-data_start[series]]
|
||||
else:
|
||||
covariates[:, 0] = cov_age[-time_len:]
|
||||
for i in range(windows_per_series[series]):
|
||||
if train:
|
||||
window_start = stride_size*i+data_start[series]
|
||||
else:
|
||||
window_start = stride_size*i
|
||||
window_end = window_start+window_size
|
||||
'''
|
||||
print("x: ", x_input[count, 1:, 0].shape)
|
||||
print("window start: ", window_start)
|
||||
print("window end: ", window_end)
|
||||
print("data: ", data.shape)
|
||||
print("d: ", data[window_start:window_end-1, series].shape)
|
||||
'''
|
||||
x_input[count, 1:, 0] = data[window_start:window_end-1, series]
|
||||
x_input[count, :, 1:1+num_covariates] = covariates[window_start:window_end, :]
|
||||
x_input[count, :, -1] = series
|
||||
label[count, :] = data[window_start:window_end, series]
|
||||
nonzero_sum = (x_input[count, 1:input_size, 0]!=0).sum()
|
||||
if nonzero_sum == 0:
|
||||
v_input[count, 0] = 0
|
||||
else:
|
||||
v_input[count, 0] = np.true_divide(x_input[count, 1:input_size, 0].sum(),nonzero_sum)+1
|
||||
x_input[count, :, 0] = x_input[count, :, 0]/v_input[count, 0]
|
||||
if train:
|
||||
label[count, :] = label[count, :]/v_input[count, 0]
|
||||
count += 1
|
||||
prefix = os.path.join(save_path, 'train_' if train else 'test_')
|
||||
np.save(prefix+'data_'+save_name, x_input)
|
||||
np.save(prefix+'v_'+save_name, v_input)
|
||||
np.save(prefix+'label_'+save_name, label)
|
||||
|
||||
def gen_covariates(times, num_covariates):
|
||||
"""Get covariates"""
|
||||
covariates = np.zeros((times.shape[0], num_covariates))
|
||||
for i, input_time in enumerate(times):
|
||||
covariates[i, 1] = input_time.weekday()
|
||||
covariates[i, 2] = input_time.hour
|
||||
covariates[i, 3] = input_time.month
|
||||
for i in range(1,num_covariates):
|
||||
covariates[:,i] = stats.zscore(covariates[:,i])
|
||||
return covariates[:, :num_covariates]
|
||||
|
||||
def visualize(data, week_start):
|
||||
x = np.arange(window_size)
|
||||
f = plt.figure()
|
||||
plt.plot(x, data[week_start:week_start+window_size], color='b')
|
||||
f.savefig("visual.png")
|
||||
plt.close()
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
global save_path
|
||||
csv_path = 'data/LD2011_2014.txt'
|
||||
save_name = 'elect'
|
||||
window_size = 192
|
||||
stride_size = 24
|
||||
num_covariates = 4
|
||||
train_start = '2011-01-01 00:00:00'
|
||||
train_end = '2014-08-31 23:00:00'
|
||||
test_start = '2014-08-25 00:00:00' #need additional 7 days as given info
|
||||
test_end = '2014-09-07 23:00:00'
|
||||
pred_days = 7
|
||||
given_days = 7
|
||||
|
||||
save_path = os.path.join('data', save_name)
|
||||
|
||||
data_frame = pd.read_csv(csv_path, sep=";", index_col=0, parse_dates=True, decimal=',')
|
||||
data_frame = data_frame.resample('1H',label = 'left',closed = 'right').sum()[train_start:test_end]
|
||||
data_frame.fillna(0, inplace=True)
|
||||
covariates = gen_covariates(data_frame[train_start:test_end].index, num_covariates)
|
||||
train_data = data_frame[train_start:train_end].values # shape: [seq_length, user_num]
|
||||
test_data = data_frame[test_start:test_end].values
|
||||
data_start = (train_data!=0).argmax(axis=0) #find first nonzero value in each time series
|
||||
total_time = data_frame.shape[0] #32304
|
||||
num_series = data_frame.shape[1] #370
|
||||
prep_data(train_data, covariates, data_start)
|
||||
prep_data(test_data, covariates, data_start, train=False)
|
||||
Vendored
-107
@@ -1,107 +0,0 @@
|
||||
from numpy.lib.npyio import save
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import os
|
||||
from tqdm import trange
|
||||
import zipfile
|
||||
|
||||
|
||||
def load_data(filedir):
|
||||
data_frame = pd.read_csv(filedir, header=0, parse_dates=True) #names=['app_name', 'zone', 'time', 'value']
|
||||
data_frame = data_frame.drop(data_frame.columns[0], axis=1)
|
||||
grouped_data = list(data_frame.groupby(["app_name", "zone"]))
|
||||
# covariates = gen_covariates(data_frame.index, 3)
|
||||
all_data = []
|
||||
for i in range(len(grouped_data)):
|
||||
single_df = grouped_data[i][1].drop(labels=['app_name', 'zone'], axis=1).sort_values(by="time", ascending=True)
|
||||
times = pd.to_datetime(single_df.time)
|
||||
single_df['weekday'] = times.dt.dayofweek / 6
|
||||
single_df['hour'] = times.dt.hour / 23
|
||||
single_df['month'] = times.dt.month / 12
|
||||
temp_data = single_df.values[:, 1:]
|
||||
if (temp_data[:, 0] == 0).sum() / len(temp_data) > 0.2:
|
||||
continue
|
||||
|
||||
all_data.append(temp_data)
|
||||
|
||||
return all_data
|
||||
|
||||
|
||||
def visualize(data, index, save_dir):
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
for i in range(index):
|
||||
x = np.arange(len(data[i]))
|
||||
f = plt.figure()
|
||||
plt.plot(x, data[i][:, 0])
|
||||
f.savefig(os.path.join(save_dir, "visual_{}.png".format(i)))
|
||||
plt.close()
|
||||
|
||||
|
||||
def split_seq(sequences, seq_length, slide_step, predict_length, save_dir):
|
||||
"""Divide the training sequence into windows"""
|
||||
train_data = []
|
||||
test_data = []
|
||||
for seq_id in trange(len(sequences)):
|
||||
split_start = 0
|
||||
single_seq = sequences[seq_id][:, 0]
|
||||
single_covariate = sequences[seq_id][:, 1:]
|
||||
windows = (len(single_seq)-seq_length+slide_step) // slide_step
|
||||
count = 0
|
||||
train_count = int(0.97 * windows)
|
||||
while len(single_seq[split_start:]) > (seq_length + predict_length):
|
||||
seq_data = single_seq[split_start:(split_start+seq_length+predict_length-1)]
|
||||
single_data = np.zeros((seq_length+predict_length-1, 5))
|
||||
single_data[:, 0] = seq_data.copy()
|
||||
single_data[:, 1:4] = single_covariate[split_start:(split_start+seq_length+predict_length-1)]
|
||||
single_data[:, -1] = seq_id
|
||||
|
||||
count += 1
|
||||
if count < train_count:
|
||||
train_data.append(single_data)
|
||||
else:
|
||||
test_data.append(single_data)
|
||||
split_start += slide_step
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
train_data = np.array(train_data, dtype=np.float32)
|
||||
train_data, v = normalize(train_data, seq_length)
|
||||
save(train_data, v, save_dir + 'train')
|
||||
test_data = np.array(test_data, dtype=np.float32)
|
||||
test_data, v = normalize(test_data, seq_length)
|
||||
save(test_data, v, save_dir + 'test')
|
||||
|
||||
|
||||
def normalize(inputs, seq_length):
|
||||
base_seq = inputs[:, :(seq_length-1), 0]
|
||||
nonzeros = (base_seq > 0).sum(1)
|
||||
v = base_seq.sum(1) / nonzeros
|
||||
v[v == 0] = 1
|
||||
inputs[:, :, 0] = inputs[:, :, 0] / v[:, None]
|
||||
|
||||
return inputs, v
|
||||
|
||||
|
||||
def save(data, v, save_dir):
|
||||
np.save(save_dir+'_data_flow.npy', data)
|
||||
np.save(save_dir+'_v_flow.npy', v)
|
||||
|
||||
|
||||
def dezip(filedir):
|
||||
zip_file = zipfile.ZipFile(filedir)
|
||||
zip_list = zip_file.namelist()
|
||||
|
||||
parent_dir = filedir.split('/')[0]
|
||||
for f in zip_list:
|
||||
zip_file.extract(f, parent_dir)
|
||||
|
||||
zip_file.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
zip_dir = 'data/app_zone_rpc_hour_encrypted.zip'
|
||||
dezip(zip_dir)
|
||||
data_dir = 'data/app_zone_rpc_hour_encrypted.csv'
|
||||
data = load_data(data_dir)
|
||||
split_seq(data, 192, 24, 24, 'data/flow/')
|
||||
Vendored
-89
@@ -1,89 +0,0 @@
|
||||
import numpy as np
|
||||
from numpy.core.defchararray import split
|
||||
import pandas as pd
|
||||
from datetime import datetime
|
||||
from scipy import stats
|
||||
import os
|
||||
|
||||
|
||||
def load_data(datadir):
|
||||
df = pd.read_csv(datadir)
|
||||
data = (df.values).transpose(1, 0)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def get_covariates(data_len, start_day):
|
||||
"""Get covariates"""
|
||||
start_timestamp = datetime.timestamp(datetime.strptime(start_day, '%Y-%m-%d %H:%M:%S'))
|
||||
timestamps = np.arange(data_len) * 3600 + start_timestamp
|
||||
timestamps = [datetime.fromtimestamp(i) for i in timestamps]
|
||||
|
||||
weekdays = stats.zscore(np.array([i.weekday() for i in timestamps]))
|
||||
hours = stats.zscore(np.array([i.hour for i in timestamps]))
|
||||
months = stats.zscore(np.array([i.month for i in timestamps]))
|
||||
|
||||
covariates = np.stack([weekdays, hours, months], axis=1)
|
||||
|
||||
return covariates
|
||||
|
||||
|
||||
def split_seq(sequences, covariates, seq_length, slide_step, predict_length, save_dir):
|
||||
"""Divide the training sequence into windows"""
|
||||
data_length = len(sequences[0])
|
||||
windows = (data_length-seq_length+slide_step) // slide_step
|
||||
train_windows = int(0.97 * windows)
|
||||
test_windows = windows - train_windows
|
||||
train_data = np.zeros((train_windows*len(sequences), seq_length+predict_length-1, 5), dtype=np.float32)
|
||||
test_data = np.zeros((test_windows*len(sequences), seq_length+predict_length-1, 5), dtype=np.float32)
|
||||
|
||||
count = 0
|
||||
split_start = 0
|
||||
seq_ids = np.arange(len(sequences))[:, None]
|
||||
end = split_start + seq_length + predict_length - 1
|
||||
while end <= data_length:
|
||||
if count < train_windows:
|
||||
train_data[count*len(sequences):(count+1)*len(sequences), :, 0] = sequences[:, split_start:end]
|
||||
train_data[count*len(sequences):(count+1)*len(sequences), :, 1:4] = covariates[split_start:end, :]
|
||||
train_data[count*len(sequences):(count+1)*len(sequences), :, -1] = seq_ids
|
||||
else:
|
||||
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, 0] = sequences[:, split_start:end]
|
||||
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, 1:4] = covariates[split_start:end, :]
|
||||
test_data[(count-train_windows)*len(sequences):(count-train_windows+1)*len(sequences), :, -1] = seq_ids
|
||||
|
||||
count += 1
|
||||
split_start += slide_step
|
||||
end = split_start + seq_length + predict_length - 1
|
||||
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
train_data, v = normalize(train_data, seq_length)
|
||||
save(train_data, v, save_dir + 'train')
|
||||
test_data, v = normalize(test_data, seq_length)
|
||||
save(test_data, v, save_dir + 'test')
|
||||
|
||||
|
||||
def normalize(inputs, seq_length):
|
||||
base_seq = inputs[:, :seq_length, 0]
|
||||
nonzeros = (base_seq > 0).sum(1)
|
||||
inputs = inputs[nonzeros > 0]
|
||||
|
||||
base_seq = inputs[:, :seq_length, 0]
|
||||
nonzeros = nonzeros[nonzeros > 0]
|
||||
v = base_seq.sum(1) / nonzeros
|
||||
v[v == 0] = 1
|
||||
inputs[:, :, 0] = inputs[:, :, 0] / v[:, None]
|
||||
|
||||
return inputs, v
|
||||
|
||||
|
||||
def save(data, v, save_dir):
|
||||
np.save(save_dir+'_data_wind.npy', data)
|
||||
np.save(save_dir+'_v_wind.npy', v)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
datadir = 'data/EMHIRESPV_TSh_CF_Country_19862015.csv'
|
||||
all_data = load_data(datadir)
|
||||
covariates = get_covariates(len(all_data[0]), '1986-01-01 00:00:00')
|
||||
split_seq(all_data, covariates, 192, 24, 24, 'data/wind/')
|
||||
Vendored
-408
@@ -1,408 +0,0 @@
|
||||
from torch.functional import align_tensors
|
||||
import torch.nn as nn
|
||||
|
||||
from torch.nn.modules.linear import Linear
|
||||
from .SubLayers import MultiHeadAttention, PositionwiseFeedForward
|
||||
import torch
|
||||
from .embed import DataEmbedding, CustomEmbedding
|
||||
import math
|
||||
|
||||
|
||||
|
||||
def get_mask(input_size, window_size, inner_size, device):
|
||||
"""Get the attention mask of PAM-Naive"""
|
||||
# Get the size of all layers
|
||||
all_size = []
|
||||
all_size.append(input_size)
|
||||
for i in range(len(window_size)):
|
||||
layer_size = math.floor(all_size[i] / window_size[i])
|
||||
all_size.append(layer_size)
|
||||
|
||||
seq_length = sum(all_size)
|
||||
mask = torch.zeros(seq_length, seq_length, device=device)
|
||||
|
||||
# get intra-scale mask
|
||||
inner_window = inner_size // 2
|
||||
for layer_idx in range(len(all_size)):
|
||||
start = sum(all_size[:layer_idx])
|
||||
for i in range(start, start + all_size[layer_idx]):
|
||||
left_side = max(i - inner_window, start)
|
||||
right_side = min(i + inner_window + 1, start + all_size[layer_idx])
|
||||
mask[i, left_side:right_side] = 1
|
||||
|
||||
# get inter-scale mask
|
||||
for layer_idx in range(1, len(all_size)):
|
||||
start = sum(all_size[:layer_idx])
|
||||
for i in range(start, start + all_size[layer_idx]):
|
||||
left_side = (start - all_size[layer_idx - 1]) + (i - start) * window_size[layer_idx - 1]
|
||||
if i == ( start + all_size[layer_idx] - 1):
|
||||
right_side = start
|
||||
else:
|
||||
right_side = (start - all_size[layer_idx - 1]) + (i - start + 1) * window_size[layer_idx - 1]
|
||||
mask[i, left_side:right_side] = 1
|
||||
mask[left_side:right_side, i] = 1
|
||||
|
||||
mask = (1 - mask).bool()
|
||||
|
||||
return mask, all_size
|
||||
|
||||
|
||||
def refer_points(all_sizes, window_size, device):
|
||||
"""Gather features from PAM's pyramid sequences"""
|
||||
input_size = all_sizes[0]
|
||||
indexes = torch.zeros(input_size, len(all_sizes), device=device)
|
||||
|
||||
for i in range(input_size):
|
||||
indexes[i][0] = i
|
||||
former_index = i
|
||||
for j in range(1, len(all_sizes)):
|
||||
start = sum(all_sizes[:j])
|
||||
inner_layer_idx = former_index - (start - all_sizes[j - 1])
|
||||
former_index = start + min(inner_layer_idx // window_size[j - 1], all_sizes[j] - 1)
|
||||
indexes[i][j] = former_index
|
||||
|
||||
indexes = indexes.unsqueeze(0).unsqueeze(3)
|
||||
|
||||
return indexes.long()
|
||||
|
||||
|
||||
def get_subsequent_mask(input_size, window_size, predict_step, truncate):
|
||||
"""Get causal attention mask for decoder."""
|
||||
if truncate:
|
||||
mask = torch.zeros(predict_step, input_size + predict_step)
|
||||
for i in range(predict_step):
|
||||
mask[i][:input_size+i+1] = 1
|
||||
mask = (1 - mask).bool().unsqueeze(0)
|
||||
else:
|
||||
all_size = []
|
||||
all_size.append(input_size)
|
||||
for i in range(len(window_size)):
|
||||
layer_size = math.floor(all_size[i] / window_size[i])
|
||||
all_size.append(layer_size)
|
||||
all_size = sum(all_size)
|
||||
mask = torch.zeros(predict_step, all_size + predict_step)
|
||||
for i in range(predict_step):
|
||||
mask[i][:all_size+i+1] = 1
|
||||
mask = (1 - mask).bool().unsqueeze(0)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def get_q_k(input_size, window_size, stride, device):
|
||||
"""
|
||||
Get the index of the key that a given query needs to attend to.
|
||||
"""
|
||||
second_length = input_size // stride
|
||||
second_last = input_size - (second_length - 1) * stride
|
||||
third_start = input_size + second_length
|
||||
third_length = second_length // stride
|
||||
third_last = second_length - (third_length - 1) * stride
|
||||
max_attn = max(second_last, third_last)
|
||||
fourth_start = third_start + third_length
|
||||
fourth_length = third_length // stride
|
||||
full_length = fourth_start + fourth_length
|
||||
fourth_last = third_length - (fourth_length - 1) * stride
|
||||
max_attn = max(third_last, fourth_last)
|
||||
|
||||
max_attn += window_size + 1
|
||||
mask = torch.zeros(full_length, max_attn, dtype=torch.int32, device=device) - 1
|
||||
|
||||
for i in range(input_size):
|
||||
mask[i, 0:window_size] = i + torch.arange(window_size) - window_size // 2
|
||||
mask[i, mask[i] > input_size - 1] = -1
|
||||
|
||||
mask[i, -1] = i // stride + input_size
|
||||
mask[i][mask[i] > third_start - 1] = third_start - 1
|
||||
for i in range(second_length):
|
||||
mask[input_size+i, 0:window_size] = input_size + i + torch.arange(window_size) - window_size // 2
|
||||
mask[input_size+i, mask[input_size+i] < input_size] = -1
|
||||
mask[input_size+i, mask[input_size+i] > third_start - 1] = -1
|
||||
|
||||
if i < second_length - 1:
|
||||
mask[input_size+i, window_size:(window_size+stride)] = torch.arange(stride) + i * stride
|
||||
else:
|
||||
mask[input_size+i, window_size:(window_size+second_last)] = torch.arange(second_last) + i * stride
|
||||
|
||||
mask[input_size+i, -1] = i // stride + third_start
|
||||
mask[input_size+i, mask[input_size+i] > fourth_start - 1] = fourth_start - 1
|
||||
for i in range(third_length):
|
||||
mask[third_start+i, 0:window_size] = third_start + i + torch.arange(window_size) - window_size // 2
|
||||
mask[third_start+i, mask[third_start+i] < third_start] = -1
|
||||
mask[third_start+i, mask[third_start+i] > fourth_start - 1] = -1
|
||||
|
||||
if i < third_length - 1:
|
||||
mask[third_start+i, window_size:(window_size+stride)] = input_size + torch.arange(stride) + i * stride
|
||||
else:
|
||||
mask[third_start+i, window_size:(window_size+third_last)] = input_size + torch.arange(third_last) + i * stride
|
||||
|
||||
mask[third_start+i, -1] = i // stride + fourth_start
|
||||
mask[third_start+i, mask[third_start+i] > full_length - 1] = full_length - 1
|
||||
for i in range(fourth_length):
|
||||
mask[fourth_start+i, 0:window_size] = fourth_start + i + torch.arange(window_size) - window_size // 2
|
||||
mask[fourth_start+i, mask[fourth_start+i] < fourth_start] = -1
|
||||
mask[fourth_start+i, mask[fourth_start+i] > full_length - 1] = -1
|
||||
|
||||
if i < fourth_length - 1:
|
||||
mask[fourth_start+i, window_size:(window_size+stride)] = third_start + torch.arange(stride) + i * stride
|
||||
else:
|
||||
mask[fourth_start+i, window_size:(window_size+fourth_last)] = third_start + torch.arange(fourth_last) + i * stride
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def get_k_q(q_k_mask):
|
||||
"""
|
||||
Get the index of the query that can attend to the given key.
|
||||
"""
|
||||
k_q_mask = q_k_mask.clone()
|
||||
for i in range(len(q_k_mask)):
|
||||
for j in range(len(q_k_mask[0])):
|
||||
if q_k_mask[i, j] >= 0:
|
||||
k_q_mask[i, j] = torch.where(q_k_mask[q_k_mask[i, j]] ==i )[0]
|
||||
|
||||
return k_q_mask
|
||||
|
||||
|
||||
class EncoderLayer(nn.Module):
|
||||
""" Compose with two layers """
|
||||
|
||||
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True, use_tvm=False, q_k_mask=None, k_q_mask=None):
|
||||
super(EncoderLayer, self).__init__()
|
||||
self.use_tvm = use_tvm
|
||||
if use_tvm:
|
||||
from .PAM_TVM import PyramidalAttention
|
||||
self.slf_attn = PyramidalAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before, q_k_mask=q_k_mask, k_q_mask=k_q_mask)
|
||||
else:
|
||||
self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)
|
||||
|
||||
self.pos_ffn = PositionwiseFeedForward(
|
||||
d_model, d_inner, dropout=dropout, normalize_before=normalize_before)
|
||||
|
||||
def forward(self, enc_input, slf_attn_mask=None):
|
||||
if self.use_tvm:
|
||||
enc_output = self.slf_attn(enc_input)
|
||||
enc_slf_attn = None
|
||||
else:
|
||||
enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)
|
||||
|
||||
enc_output = self.pos_ffn(enc_output)
|
||||
|
||||
return enc_output, enc_slf_attn
|
||||
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
""" Compose with two layers """
|
||||
|
||||
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1, normalize_before=True):
|
||||
super(DecoderLayer, self).__init__()
|
||||
self.slf_attn = MultiHeadAttention(
|
||||
n_head, d_model, d_k, d_v, dropout=dropout, normalize_before=normalize_before)
|
||||
self.pos_ffn = PositionwiseFeedForward(
|
||||
d_model, d_inner, dropout=dropout, normalize_before=normalize_before)
|
||||
|
||||
def forward(self, Q, K, V, slf_attn_mask=None):
|
||||
enc_output, enc_slf_attn = self.slf_attn(
|
||||
Q, K, V, mask=slf_attn_mask)
|
||||
|
||||
enc_output = self.pos_ffn(enc_output)
|
||||
|
||||
return enc_output, enc_slf_attn
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
def __init__(self, c_in, window_size):
|
||||
super(ConvLayer, self).__init__()
|
||||
self.downConv = nn.Conv1d(in_channels=c_in,
|
||||
out_channels=c_in,
|
||||
kernel_size=window_size,
|
||||
stride=window_size)
|
||||
self.norm = nn.BatchNorm1d(c_in)
|
||||
self.activation = nn.ELU()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.downConv(x)
|
||||
x = self.norm(x)
|
||||
x = self.activation(x)
|
||||
return x
|
||||
|
||||
|
||||
class Conv_Construct(nn.Module):
|
||||
"""Convolution CSCM"""
|
||||
def __init__(self, d_model, window_size, d_inner):
|
||||
super(Conv_Construct, self).__init__()
|
||||
if not isinstance(window_size, list):
|
||||
self.conv_layers = nn.ModuleList([
|
||||
ConvLayer(d_model, window_size),
|
||||
ConvLayer(d_model, window_size),
|
||||
ConvLayer(d_model, window_size)
|
||||
])
|
||||
else:
|
||||
self.conv_layers = nn.ModuleList([
|
||||
ConvLayer(d_model, window_size[0]),
|
||||
ConvLayer(d_model, window_size[1]),
|
||||
ConvLayer(d_model, window_size[2])
|
||||
])
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, enc_input):
|
||||
all_inputs = []
|
||||
enc_input = enc_input.permute(0, 2, 1)
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
for i in range(len(self.conv_layers)):
|
||||
enc_input = self.conv_layers[i](enc_input)
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
|
||||
all_inputs = self.norm(all_inputs)
|
||||
|
||||
return all_inputs
|
||||
|
||||
|
||||
class Bottleneck_Construct(nn.Module):
|
||||
"""Bottleneck convolution CSCM"""
|
||||
def __init__(self, d_model, window_size, d_inner):
|
||||
super(Bottleneck_Construct, self).__init__()
|
||||
if not isinstance(window_size, list):
|
||||
self.conv_layers = nn.ModuleList([
|
||||
ConvLayer(d_inner, window_size),
|
||||
ConvLayer(d_inner, window_size),
|
||||
ConvLayer(d_inner, window_size)
|
||||
])
|
||||
else:
|
||||
self.conv_layers = []
|
||||
for i in range(len(window_size)):
|
||||
self.conv_layers.append(ConvLayer(d_inner, window_size[i]))
|
||||
self.conv_layers = nn.ModuleList(self.conv_layers)
|
||||
self.up = Linear(d_inner, d_model)
|
||||
self.down = Linear(d_model, d_inner)
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, enc_input):
|
||||
|
||||
temp_input = self.down(enc_input).permute(0, 2, 1)
|
||||
all_inputs = []
|
||||
for i in range(len(self.conv_layers)):
|
||||
temp_input = self.conv_layers[i](temp_input)
|
||||
all_inputs.append(temp_input)
|
||||
|
||||
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
|
||||
all_inputs = self.up(all_inputs)
|
||||
all_inputs = torch.cat([enc_input, all_inputs], dim=1)
|
||||
|
||||
all_inputs = self.norm(all_inputs)
|
||||
|
||||
return all_inputs
|
||||
|
||||
|
||||
class MaxPooling_Construct(nn.Module):
|
||||
"""Max pooling CSCM"""
|
||||
def __init__(self, d_model, window_size, d_inner):
|
||||
super(MaxPooling_Construct, self).__init__()
|
||||
if not isinstance(window_size, list):
|
||||
self.pooling_layers = nn.ModuleList([
|
||||
nn.MaxPool1d(kernel_size=window_size),
|
||||
nn.MaxPool1d(kernel_size=window_size),
|
||||
nn.MaxPool1d(kernel_size=window_size)
|
||||
])
|
||||
else:
|
||||
self.pooling_layers = nn.ModuleList([
|
||||
nn.MaxPool1d(kernel_size=window_size[0]),
|
||||
nn.MaxPool1d(kernel_size=window_size[1]),
|
||||
nn.MaxPool1d(kernel_size=window_size[2])
|
||||
])
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, enc_input):
|
||||
all_inputs = []
|
||||
enc_input = enc_input.transpose(1, 2).contiguous()
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
for layer in self.pooling_layers:
|
||||
enc_input = layer(enc_input)
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
|
||||
all_inputs = self.norm(all_inputs)
|
||||
|
||||
return all_inputs
|
||||
|
||||
|
||||
class AvgPooling_Construct(nn.Module):
|
||||
"""Average pooling CSCM"""
|
||||
def __init__(self, d_model, window_size, d_inner):
|
||||
super(AvgPooling_Construct, self).__init__()
|
||||
if not isinstance(window_size, list):
|
||||
self.pooling_layers = nn.ModuleList([
|
||||
nn.AvgPool1d(kernel_size=window_size),
|
||||
nn.AvgPool1d(kernel_size=window_size),
|
||||
nn.AvgPool1d(kernel_size=window_size)
|
||||
])
|
||||
else:
|
||||
self.pooling_layers = nn.ModuleList([
|
||||
nn.AvgPool1d(kernel_size=window_size[0]),
|
||||
nn.AvgPool1d(kernel_size=window_size[1]),
|
||||
nn.AvgPool1d(kernel_size=window_size[2])
|
||||
])
|
||||
self.norm = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, enc_input):
|
||||
all_inputs = []
|
||||
enc_input = enc_input.transpose(1, 2).contiguous()
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
for layer in self.pooling_layers:
|
||||
enc_input = layer(enc_input)
|
||||
all_inputs.append(enc_input)
|
||||
|
||||
all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)
|
||||
all_inputs = self.norm(all_inputs)
|
||||
|
||||
return all_inputs
|
||||
|
||||
|
||||
class Predictor(nn.Module):
|
||||
|
||||
def __init__(self, dim, num_types):
|
||||
super().__init__()
|
||||
|
||||
self.linear = nn.Linear(dim, num_types, bias=False)
|
||||
nn.init.xavier_normal_(self.linear.weight)
|
||||
|
||||
def forward(self, data):
|
||||
out = self.linear(data)
|
||||
out = out
|
||||
return out
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
""" A encoder model with self attention mechanism. """
|
||||
|
||||
def __init__(self, opt, mask):
|
||||
super().__init__()
|
||||
|
||||
self.model_type = opt.model
|
||||
self.mask = mask
|
||||
|
||||
self.layers = nn.ModuleList([
|
||||
DecoderLayer(opt.d_model, opt.d_inner_hid, opt.n_head, opt.d_k, opt.d_v, dropout=opt.dropout, \
|
||||
normalize_before=False),
|
||||
DecoderLayer(opt.d_model, opt.d_inner_hid, opt.n_head, opt.d_k, opt.d_v, dropout=opt.dropout, \
|
||||
normalize_before=False)
|
||||
])
|
||||
|
||||
if opt.embed_type == 'CustomEmbedding':
|
||||
self.dec_embedding = CustomEmbedding(opt.enc_in, opt.d_model, opt.covariate_size, opt.seq_num, opt.dropout)
|
||||
else:
|
||||
self.dec_embedding = DataEmbedding(opt.enc_in, opt.d_model, opt.dropout)
|
||||
|
||||
def forward(self, x_dec, x_mark_dec, refer):
|
||||
dec_enc = self.dec_embedding(x_dec, x_mark_dec)
|
||||
|
||||
dec_enc, _ = self.layers[0](dec_enc, refer, refer)
|
||||
refer_enc = torch.cat([refer, dec_enc], dim=1)
|
||||
mask = self.mask.repeat(len(dec_enc), 1, 1).to(dec_enc.device)
|
||||
dec_enc, _ = self.layers[1](dec_enc, refer_enc, refer_enc, slf_attn_mask=mask)
|
||||
|
||||
return dec_enc
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user