diff --git a/code/new_realtime/.RData b/code/new_realtime/.RData new file mode 100644 index 0000000..dbeff96 Binary files /dev/null and b/code/new_realtime/.RData differ diff --git a/code/new_realtime/.Rhistory b/code/new_realtime/.Rhistory new file mode 100644 index 0000000..597ab21 --- /dev/null +++ b/code/new_realtime/.Rhistory @@ -0,0 +1,11 @@ +install.package("cpm") +install.packages("cpm") +install.packages("cpm") +install.packages("signal", "zoo", "pracma") +install.packages("signal", "zoo", "pracma") +install.packages("signal", "pracma") +install.packages("signal") +install.packages("signal") +install.packages("signal", "zoo", "pracma") +install.packages("zoo", "pracma") +install.packages("pracma") diff --git a/code/new_realtime/.env b/code/new_realtime/.env new file mode 100644 index 0000000..92bad2e --- /dev/null +++ b/code/new_realtime/.env @@ -0,0 +1,3 @@ +LOG_DIR=/home/alex/projects/datascience-analysis/notebooks/logs +LMDB_ROOT_DIR=/home/alex/projects/datascience-analysis/notebooks/lmdb_datasets +RESULTS_ROOT_DIR=/home/alex/projects/datascience-analysis/notebooks/results \ No newline at end of file diff --git a/code/new_realtime/__pycache__/experiment_setup.cpython-310.pyc b/code/new_realtime/__pycache__/experiment_setup.cpython-310.pyc new file mode 100644 index 0000000..8c3a511 Binary files /dev/null and b/code/new_realtime/__pycache__/experiment_setup.cpython-310.pyc differ diff --git a/code/new_realtime/__pycache__/run_config.cpython-310.pyc b/code/new_realtime/__pycache__/run_config.cpython-310.pyc new file mode 100644 index 0000000..6fc356f Binary files /dev/null and b/code/new_realtime/__pycache__/run_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__init__.py b/code/new_realtime/configs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/code/new_realtime/configs/__pycache__/__init__.cpython-310.pyc b/code/new_realtime/configs/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..dd5b443 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/__init__.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/cnn_run_config.cpython-310.pyc b/code/new_realtime/configs/__pycache__/cnn_run_config.cpython-310.pyc new file mode 100644 index 0000000..737daf5 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/cnn_run_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/combined.cpython-310.pyc b/code/new_realtime/configs/__pycache__/combined.cpython-310.pyc new file mode 100644 index 0000000..68b4564 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/combined.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/feature_config.cpython-310.pyc b/code/new_realtime/configs/__pycache__/feature_config.cpython-310.pyc new file mode 100644 index 0000000..6f27713 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/feature_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/lstm_run_config.cpython-310.pyc b/code/new_realtime/configs/__pycache__/lstm_run_config.cpython-310.pyc new file mode 100644 index 0000000..fe7a01b Binary files /dev/null and b/code/new_realtime/configs/__pycache__/lstm_run_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/patch_run_config.cpython-310.pyc b/code/new_realtime/configs/__pycache__/patch_run_config.cpython-310.pyc new file mode 100644 index 0000000..dd0f7c0 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/patch_run_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/transformer.cpython-310.pyc b/code/new_realtime/configs/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000..56c41d4 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/transformer.cpython-310.pyc differ diff --git a/code/new_realtime/configs/__pycache__/transformer_run_config.cpython-310.pyc b/code/new_realtime/configs/__pycache__/transformer_run_config.cpython-310.pyc new file mode 100644 index 0000000..5f983d6 Binary files /dev/null and b/code/new_realtime/configs/__pycache__/transformer_run_config.cpython-310.pyc differ diff --git a/code/new_realtime/configs/cnn_run_config.py b/code/new_realtime/configs/cnn_run_config.py new file mode 100644 index 0000000..88aafc5 --- /dev/null +++ b/code/new_realtime/configs/cnn_run_config.py @@ -0,0 +1,127 @@ +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, + ) + ] +} diff --git a/code/new_realtime/configs/combined.py b/code/new_realtime/configs/combined.py new file mode 100644 index 0000000..f0670a3 --- /dev/null +++ b/code/new_realtime/configs/combined.py @@ -0,0 +1,132 @@ +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": 256, + "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, + }, + }, + { + "model_parameters": { + "cnn_channels": 64, + "kernel_size": 3, + "embed_dim": 64, + "num_enc_layers": 2, + "num_heads": 2, + }, + }, + { + "batch_size": 128, + "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, + }, + }, + ] + ) +} diff --git a/code/new_realtime/configs/feature_config.py b/code/new_realtime/configs/feature_config.py new file mode 100644 index 0000000..245cd9c --- /dev/null +++ b/code/new_realtime/configs/feature_config.py @@ -0,0 +1,167 @@ +import os +from functools import partial +import math + +import torch +from torch import nn +import numpy as np + +from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler +from torch.optim.lr_scheduler import OneCycleLR +from torch.optim import AdamW + +from data_analysis.models.ov_detection.config import model_config +from utils.feature_functions import * +from utils.loss_functions import weighted_bce_loss_fn +from utils.training_utils import collate + +from vsm_datascience_common import constants + +feature_config = { + "feature_set_name": "full_feature_set", + "filter_criteria": { + "measurements.length": { + "$gt": constants.MEASUREMENTS_PER_DAY * 10, + "$lt": constants.MEASUREMENTS_PER_DAY * 150 + }, + "ends_at": { + "$exists": True, + "$ne": None + }, + "$and": [ + {"measurements.values": {"$not": {"$elemMatch": {"$lt": 35}}}}, + {"measurements.values": {"$not": {"$elemMatch": {"$gt": 43}}}} + ] + }, + "augmentation": { + "use_augmentation": True, + "max_lookback": 4, + }, + "feature_sets": [ + "static_categorical_features", + "static_continuous_features", + "known_categorical_features", + "known_continuous_features", + "observed_categorical_features", + "observed_continuous_features", + "target_features" + ], + "static_categorical_features": [ + ], + "static_continuous_features": [ + { + "name": "average_cycle_length", + "fn": get_cycle_length_stats, + "scaler": RobustScaler + }, + { + "name": "average_ovulation_day", + "fn": get_average_ovulation_day, + "scaler": RobustScaler + }, + { + "name": "ovulation_std", + "fn": get_ovulation_std, + "scaler": StandardScaler + }, + { + "name": "biphasic_fraction", + "fn": get_biphasic_fraction, + "scaler": MinMaxScaler + }, + { + "name": "num_cycles", + "fn": get_num_cycles, + "scaler": RobustScaler + }, + { + "name": "temperature_averages", + "fn": get_average_temperatures, + "scaler": StandardScaler + } + ], + "known_categorical_features": [ + ], + "known_continuous_features": [ + { + "name": "hours_from_start", + "fn": partial(get_hours_from_start, shift=0), + "scaler": MinMaxScaler, + "accumulation_fn": np.max + }, + { + "name": "hour_of_day", + "fn": partial(get_hour_of_day_encoded, shift=0), + "scaler": None, + "accumulation_fn": np.max + }, + { + "name": "day_of_week", + "fn": partial(get_day_of_week_encoded, shift=0), + "scaler": None, + "accumulation_fn": np.max + }, + { + "name": "month_of_year", + "fn": partial(get_month_of_year_encoded, shift=0), + "scaler": None, + "accumulation_fn": np.max + } + ], + "observed_categorical_features": [ + ], + "observed_continuous_features": [ + { + "name": "temperature", + "fn": partial(get_temperature, shift=0), + "scaler": StandardScaler + }, + { + "name": "rolling_average_temperature", + "fn": partial(get_rolling_average_with_padding, shift=0), + "scaler": StandardScaler + }, + { + "name": "rolling_window_temperature_min", + "fn": partial(get_window_fn, window_size=constants.MEASUREMENTS_PER_DAY, fn=partial(np.min, axis=1)), + "scaler": StandardScaler + }, + { + "name": "rolling_window_temperature_max", + "fn": partial(get_window_fn, window_size=constants.MEASUREMENTS_PER_DAY, fn=partial(np.max, axis=1)), + "scaler": StandardScaler + }, + ], + "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": "days_relative_to_ov", + "fn": get_days_relative_to_ov, + "scaler": RobustScaler, + "accumulation_fn": np.max + }, + { + "name": "ov_day", + "fn": get_ov_day, + "scaler": RobustScaler, + "accumulation_fn": np.max + }, + { + "name": "is_biphasic", + "fn": get_is_biphasic, + "scaler": MinMaxScaler, + "accumulation_fn": np.max + } + ] +} diff --git a/code/new_realtime/configs/lstm_run_config.py b/code/new_realtime/configs/lstm_run_config.py new file mode 100644 index 0000000..7aed0b1 --- /dev/null +++ b/code/new_realtime/configs/lstm_run_config.py @@ -0,0 +1,127 @@ +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 + ) + ] +} diff --git a/code/new_realtime/configs/patch_run_config.py b/code/new_realtime/configs/patch_run_config.py new file mode 100644 index 0000000..3448721 --- /dev/null +++ b/code/new_realtime/configs/patch_run_config.py @@ -0,0 +1,109 @@ +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, + } + + } + ] +} diff --git a/code/new_realtime/configs/transformer_run_config.py b/code/new_realtime/configs/transformer_run_config.py new file mode 100644 index 0000000..958064c --- /dev/null +++ b/code/new_realtime/configs/transformer_run_config.py @@ -0,0 +1,123 @@ +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, + ) + ] +} diff --git a/code/new_realtime/dataset_wrapper.py b/code/new_realtime/dataset_wrapper.py new file mode 100644 index 0000000..a3c8cfa --- /dev/null +++ b/code/new_realtime/dataset_wrapper.py @@ -0,0 +1,156 @@ +import sys +import argparse +from concurrent.futures import ProcessPoolExecutor +from functools import partial +import logging + +import lmdb +import dotenv +from tqdm import tqdm + +dotenv.load_dotenv() + +from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id +from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection + +from utils.dataset_creation import get_features, train_scalers, save_scalers, scale_item +from utils.lmdb_utils import save_to_lmdb, load_from_lmdb +from utils.utils import get_variable_from_module + +MAX_LMDB_SIZE_IN_MB = 200_000 + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) +logger.addHandler(logging.StreamHandler(sys.stdout)) + + +def process_id_wrapper(cycle_id: str, feature_config: dict, env: lmdb.Environment): + try: + cycle = get_cycle_by_id(cycle_id) + features = get_features(cycle, feature_config) + save_to_lmdb(env, key=str(cycle_id), dataset=features) + except: + pass + + +def scaling_wrapper(key_batch: str, + scalers: dict, + max_lmdb_size_in_mb: int, + lmdb_env_dir: str): + env = lmdb.open(lmdb_env_dir, readonly=True, lock=False) + scaled_items = list() + for key in key_batch: + item = load_from_lmdb(env, key) + scaled_item = scale_item(item, scalers) + scaled_items.append(scaled_item) + env.close() + env = lmdb.open(lmdb_env_dir, map_size=max_lmdb_size_in_mb * 1024 * 1024) + for key, scaled_item in zip(key_batch, scaled_items): + save_to_lmdb(env, key=key, dataset=scaled_item) + env.close() + + +def create_dataset(model_configuration: dict, + lmdb_root_dir: str, + max_lmdb_size_in_mb: int, + max_workers: int) -> None: + feature_config = model_configuration["feature_config"] + + logger.info(f"Creating dataset for feature set {feature_config['feature_set_name']}") + + # fetch valid cycle ids from database + valid_cycle_ids = [x["_id"] for x in get_cycles_collection().aggregate( + feature_config["filter_criteria_pipeline"] + [ + { + "$project": { + "_id": 1 + } + } + ] + )] + + logger.info(f"Fetched {len(valid_cycle_ids)} valid cycle ids from database") + + env = lmdb.open(f"{lmdb_root_dir}/{feature_config['feature_set_name']}", map_size=max_lmdb_size_in_mb * 1024 * 1024) + + # create features for items + logger.info(f"Creating features for {len(valid_cycle_ids)} cycles") + with ProcessPoolExecutor(max_workers=max_workers) as executor: + list(tqdm(executor.map(partial(process_id_wrapper, feature_config=feature_config, env=env), valid_cycle_ids), + total=len(valid_cycle_ids))) + + # convert ids (here bson objectids) to keys for use in lmdb + keys = [str(cycle_id) for cycle_id in valid_cycle_ids] + + # train scalers for featues + logger.info("Training scalers for features") + all_scalers = dict() + sample = load_from_lmdb(env, keys[0]) + for feature_set in tqdm(feature_config["feature_sets"]): + for feature in feature_config[feature_set]: + print(f"Training scalers for feature {feature['name']} in feature set {feature_set}") + scaler = train_scalers(feature_set, feature["name"], feature["scaler"], sample, env) + if feature_set not in all_scalers: + all_scalers[feature_set] = dict() + all_scalers[feature_set] = all_scalers[feature_set] | scaler + + # save scalers + logger.info("Saving scalers to disk") + scaler_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}/scalers" + save_scalers(all_scalers, scaler_dir) + + # creat scaling batches for less burden on lmdb + batch_size = 1_000 + key_batches = [keys[i:i + batch_size] for i in range(0, len(keys), batch_size)] + + # scale items + logger.info("Scaling items") + with ProcessPoolExecutor(max_workers=max_workers) as executor: + list(tqdm(executor.map(partial(scaling_wrapper, scalers=all_scalers, + max_lmdb_size_in_mb=max_lmdb_size_in_mb, + lmdb_env_dir=f"{lmdb_root_dir}/{feature_config['feature_set_name']}"), + key_batches), + total=len(key_batches))) + + logger.info(f"Dataset creation finished. LMDB saved in {lmdb_root_dir}/{feature_config['feature_set_name']}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Dataset Creation Wrapper") + parser.add_argument("--model_config_module", + type=str, + required=True, + help="Path to the model config module") + parser.add_argument("--model_config_variable", + type=str, + required=False, + default="model_configuration", + help="Name of the model config variable, default is 'model_configuration'") + parser.add_argument("--lmdb_dir", + type=str, + required=False, + default="./lmdb_datasets", + help="Path to the lmdb directory") + parser.add_argument("--lmdb_size", + type=int, + required=False, + default=MAX_LMDB_SIZE_IN_MB, + help="Size of the lmdb in MB, default is 200_000") + parser.add_argument("--max_workers", + type=int, + required=False, + default=None, + help="Number of workers for multiprocessing, default is None (use all available cores)") + + args = parser.parse_args() + + # import and load the model config + model_configuration = get_variable_from_module( + module_path=args.model_config_module, + variable_name=args.model_config_variable + ) + + create_dataset(model_configuration, + lmdb_root_dir=args.lmdb_dir, + max_lmdb_size_in_mb=args.lmdb_size, + max_workers=args.max_workers) diff --git a/code/new_realtime/experiment_setup.py b/code/new_realtime/experiment_setup.py new file mode 100644 index 0000000..6c7f6fb --- /dev/null +++ b/code/new_realtime/experiment_setup.py @@ -0,0 +1,48 @@ +from functools import partial + +import numpy as np +import sklearn + +from utils.evaluation import * + + +def get_eval_functions(model_configuration: dict): + """ + Returns the evaluation functions for the model + :return: list of evaluation functions + """ + 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), + "accumulation_fn": np.mean, + } + ] + return eval_functions diff --git a/code/new_realtime/models/__init__.py b/code/new_realtime/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/code/new_realtime/models/__pycache__/__init__.cpython-310.pyc b/code/new_realtime/models/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..95f01f6 Binary files /dev/null and b/code/new_realtime/models/__pycache__/__init__.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/cnn.cpython-310.pyc b/code/new_realtime/models/__pycache__/cnn.cpython-310.pyc new file mode 100644 index 0000000..9213b99 Binary files /dev/null and b/code/new_realtime/models/__pycache__/cnn.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/collation.cpython-310.pyc b/code/new_realtime/models/__pycache__/collation.cpython-310.pyc new file mode 100644 index 0000000..7e73894 Binary files /dev/null and b/code/new_realtime/models/__pycache__/collation.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/lstm.cpython-310.pyc b/code/new_realtime/models/__pycache__/lstm.cpython-310.pyc new file mode 100644 index 0000000..fea9140 Binary files /dev/null and b/code/new_realtime/models/__pycache__/lstm.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/tft_utils.cpython-310.pyc b/code/new_realtime/models/__pycache__/tft_utils.cpython-310.pyc new file mode 100644 index 0000000..034fd60 Binary files /dev/null and b/code/new_realtime/models/__pycache__/tft_utils.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/transformer.cpython-310.pyc b/code/new_realtime/models/__pycache__/transformer.cpython-310.pyc new file mode 100644 index 0000000..6590bb6 Binary files /dev/null and b/code/new_realtime/models/__pycache__/transformer.cpython-310.pyc differ diff --git a/code/new_realtime/models/__pycache__/utils.cpython-310.pyc b/code/new_realtime/models/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000..8c2eb17 Binary files /dev/null and b/code/new_realtime/models/__pycache__/utils.cpython-310.pyc differ diff --git a/code/new_realtime/models/cnn.py b/code/new_realtime/models/cnn.py new file mode 100644 index 0000000..1f50f21 --- /dev/null +++ b/code/new_realtime/models/cnn.py @@ -0,0 +1,58 @@ +import math + +import torch +from torch import nn + + +class CNNTransformer(nn.Module): + def __init__(self, + 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), + nn.ReLU(), + nn.Conv1d(cnn_channels, embed_dim, kernel_size=kernel_size, padding=1), + nn.BatchNorm1d(embed_dim), + nn.ReLU() + ) + + # 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) + self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers) + self.pool = nn.AdaptiveAvgPool1d(1) + self.head = nn.Linear(embed_dim, output_dim) + + def _get_sinusoidal_embedding(self, seq_len, embed_dim): + position = torch.arange(0, seq_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim)) + pe = torch.zeros(seq_len, embed_dim) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + 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 + + # Add positional embedding + pos_embed = self.pos_embed[:, :cnn_out.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 + + # Apply Transformer encoder + enc = self.encoder(cnn_out) # (S, B, E) + pooled = enc.mean(0) # (B, E) + return self.head(pooled) diff --git a/code/new_realtime/models/collation.py b/code/new_realtime/models/collation.py new file mode 100644 index 0000000..77e3e24 --- /dev/null +++ b/code/new_realtime/models/collation.py @@ -0,0 +1,94 @@ +import numpy as np +import torch + + +def flat_x_y_collate(batch, *args, **kwargs) -> tuple: + """ + Collate function for a classical model with flat x and y feature vectors + """ + # Flatten each item in the batch + flattened_x = list() + flattened_y = list() + for i in range(len(batch)): + item = batch[i] + flattened_x_item = list() + flattened_y_item = list() + for j, (category, windows) in enumerate(item.items()): + for feature in windows.T: + if category == "target": + flattened_y_item.append(feature) + else: + flattened_x_item.append(feature) + # create numpy arrays with t_1_f1, t_1_f2, t_2_f1, t_2_f2 and so on + reordered_x = np.empty((len(flattened_x_item) * len(flattened_x_item[0]))) + reordered_y = np.empty((len(flattened_y_item) * len(flattened_y_item[0]))) + for j in range(len(flattened_x_item)): + for k in range(len(flattened_x_item[j])): + reordered_x[j + k * len(flattened_x_item)] = flattened_x_item[j][k] + + for j in range(len(flattened_y_item)): + for k in range(len(flattened_y_item[j])): + reordered_y[j + k * len(flattened_y_item)] = flattened_y_item[j][k] + + # append to the list + flattened_x.append(reordered_x) + flattened_y.append(reordered_y) + + return np.array(flattened_x), np.array(flattened_y) + + +def simple_x_y_collate(batch): + """ + Collate function for LSTM model + """ + collated_x = list() + collated_y = list() + for item in batch: + current_x = list() + current_y = list() + for category, windows in item.items(): + if category == "target_features": + current_y.append(windows) + else: + current_x.append(windows) + + collated_x.append(np.concatenate(current_x, axis=1)) + collated_y.append(np.concatenate(current_y, axis=1)) + + return torch.tensor(collated_x, dtype=torch.float32), torch.tensor(collated_y, dtype=torch.float32) + + +def collate_with_padding(batch, + padding_value: float = 0.0, ): + """ + Collate the batch with padding + """ + + # create simple lists for x and y + collated_x = list() + collated_y = list() + for item in batch: + current_x = list() + current_y = list() + for category, windows in item.items(): + if category == "target_features": + current_y.append(windows) + else: + current_x.append(windows) + + collated_x.append(np.concatenate(current_x, axis=1)) + collated_y.append(np.concatenate(current_y, axis=1)) + + # get the max length of the x and y + max_x_length = max([x.shape[0] for x in collated_x]) + max_y_length = max([y.shape[0] for y in collated_y]) + # pad the x and y + padded_x = list() + padded_y = list() + for x, y in zip(collated_x, collated_y): + padded_x.append( + np.pad(x, ((max_x_length - x.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value)) + padded_y.append( + np.pad(y, ((max_y_length - y.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value)) + + return torch.tensor(padded_x, dtype=torch.float32), torch.tensor(padded_y, dtype=torch.float32) diff --git a/code/new_realtime/models/lstm.py b/code/new_realtime/models/lstm.py new file mode 100644 index 0000000..6181a3b --- /dev/null +++ b/code/new_realtime/models/lstm.py @@ -0,0 +1,46 @@ +from torch import nn + + +class LSTMModel(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): + 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, + hidden_size=lstm_hidden_size, + num_layers=num_layers, + batch_first=True, + ) + 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.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 + # x: (batch_size, lstm_hidden_size) + x = self.head(x) + # x: (batch_size, output_size) + return x diff --git a/code/new_realtime/models/tft_utils.py b/code/new_realtime/models/tft_utils.py new file mode 100644 index 0000000..d16e933 --- /dev/null +++ b/code/new_realtime/models/tft_utils.py @@ -0,0 +1,155 @@ +import datetime +import os + +import torch +from torch import nn +from torch.optim import AdamW +from torch.optim.lr_scheduler import OneCycleLR +from torch.utils.data import IterableDataset, DataLoader + +from models.third_party.tft_model import TemporalFusionTransformer + + +def get_tft_model(model_configuration: dict, + sample_item: dict, + device: str) -> nn.Module: + config_class = create_config_class(model_configuration, sample_item) + model = TemporalFusionTransformer(config_class) + model.to(device) + return model + + +def create_training_state(model_configuration: dict, + training_configuration: dict, + train_dataset: IterableDataset | DataLoader, + device: str) -> dict: + training_state = dict() + + sample_item = next(iter(train_dataset)) + + if model_configuration["model_type"] == "TemporalFusionTransformer": + training_state["model"] = get_tft_model(model_configuration, + sample_item, + device) + else: + raise NotImplementedError(f"Model type {model_configuration['type']} not implemented") + + training_state["optimizer"] = AdamW(training_state["model"].parameters(), + lr=training_configuration["learning_rate"]) + training_state["scheduler"] = OneCycleLR(training_state["optimizer"], + max_lr=training_configuration["learning_rate"], + total_steps=len(train_dataset) * + training_configuration[ + "epochs"]) + training_state["current_epoch"] = 1 + return training_state + + +def get_checkpoints(model_configuration: dict, training_configuration: dict) -> list: + checkpoints_dir = f"{model_configuration['model_dir']}/trainings/{training_configuration['id']}/checkpoints" + if os.path.exists(checkpoints_dir): + checkpoints = [os.path.join(checkpoints_dir, f) for f in os.listdir(checkpoints_dir) if + f.endswith('.pt') and "checkpoint" in f] + checkpoints.sort(key=os.path.getmtime, reverse=False) + return checkpoints + else: + return [] + + +def load_checkpoint(checkpoint_path: str, + training_configuration: dict, + model_configuration: dict, + test_data_loader: IterableDataset | DataLoader, + device: str) -> tuple: + checkpoint = torch.load(checkpoint_path) + training_state = torch.load(checkpoint_path) + + sample_item = next(iter(test_data_loader)) + + # load model + if model_configuration["model_type"] == "TemporalFusionTransformer": + model = get_tft_model(model_configuration, + sample_item, + device) + model_configuration["model"] = model + else: + raise NotImplementedError(f"Model type {model_configuration['model_type']} not implemented") + + # load optimizer + optimizer = AdamW(model.parameters(), lr=training_configuration["learning_rate"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + training_state["optimizer"] = optimizer + + # load scheduler + scheduler = OneCycleLR(optimizer, + max_lr=training_configuration["learning_rate"], + total_steps=len(test_data_loader) * training_configuration["epochs"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + training_state["scheduler"] = scheduler + + return training_state + + +def save_checkpoint(training_config: dict, + training_state: dict, + model_config: dict) -> None: + checkpoints_dir = f"{model_config['model_dir']}/trainings/{training_config['id']}/checkpoints" + checkpoint_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + checkpoint_path = os.path.join(checkpoints_dir, f"checkpoint_{checkpoint_id}.pt") + if not os.path.exists(checkpoints_dir): + os.makedirs(checkpoints_dir) + + config_to_save = training_state.copy() + + # replace training parts with their state dicts + config_to_save["model"] = config_to_save["model"].state_dict() + config_to_save["optimizer"] = config_to_save["optimizer"].state_dict() + config_to_save["scheduler"] = config_to_save["scheduler"].state_dict() + + # save training state + torch.save(config_to_save, checkpoint_path) + + +def create_config_class(config: dict, sample_batch: dict) -> object: + class ConfigClass: + def __init__(self): + # Feature sizes + self.static_categorical_inp_lens = [] + self.temporal_known_categorical_inp_lens = [] + self.temporal_observed_categorical_inp_lens = [] + + model_parameters = config["model_parameters"] + + self.example_length = model_parameters["encoder_length"] + model_parameters["decoder_length"] + self.encoder_length = model_parameters["encoder_length"] + + self.n_head = model_parameters["attention_heads"] + self.hidden_size = model_parameters["state_size"] + self.dropout = model_parameters["dropout"] + self.attn_dropout = model_parameters["attention_dropout"] + self.quantiles = model_parameters["output_quantiles"] + self.use_past_targets = model_parameters["use_past_targets"] + + #### Derived variables #### + self.temporal_known_continuous_inp_size = sample_batch["k_cont"].shape[2] + self.temporal_observed_continuous_inp_size = sample_batch["o_cont"].shape[2] + self.temporal_target_size = sample_batch["target"].shape[2] + self.static_continuous_inp_size = sample_batch["s_cont"].shape[2] + + self.num_static_vars = self.static_continuous_inp_size + len(self.static_categorical_inp_lens) + self.num_future_vars = self.temporal_known_continuous_inp_size + len( + self.temporal_known_categorical_inp_lens) + if self.use_past_targets: + self.num_historic_vars = self.num_future_vars + self.temporal_observed_continuous_inp_size + self.temporal_target_size + len( + self.temporal_observed_categorical_inp_lens) + else: + self.num_historic_vars = self.num_future_vars + self.temporal_observed_continuous_inp_size + len( + self.temporal_observed_categorical_inp_lens) + # self.num_historic_vars = sum([self.num_future_vars, + # self.temporal_observed_continuous_inp_size, + # self.temporal_target_size, + # len(self.temporal_observed_categorical_inp_lens), + # ]) + self.target_size = self.temporal_target_size + + return ConfigClass() diff --git a/code/new_realtime/models/third_party/__pycache__/__init__.cpython-310.pyc b/code/new_realtime/models/third_party/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..9719d0f Binary files /dev/null and b/code/new_realtime/models/third_party/__pycache__/__init__.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/__pycache__/tft_model.cpython-310.pyc b/code/new_realtime/models/third_party/__pycache__/tft_model.cpython-310.pyc new file mode 100644 index 0000000..3ae70d4 Binary files /dev/null and b/code/new_realtime/models/third_party/__pycache__/tft_model.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst/__init__.py b/code/new_realtime/models/third_party/patch_tst/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/code/new_realtime/models/third_party/patch_tst/__pycache__/__init__.cpython-310.pyc b/code/new_realtime/models/third_party/patch_tst/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..250d26f Binary files /dev/null and b/code/new_realtime/models/third_party/patch_tst/__pycache__/__init__.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst/layers/AutoCorrelation.py b/code/new_realtime/models/third_party/patch_tst/layers/AutoCorrelation.py new file mode 100644 index 0000000..a6fb63c --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/AutoCorrelation.py @@ -0,0 +1,164 @@ +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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/Autoformer_EncDec.py b/code/new_realtime/models/third_party/patch_tst/layers/Autoformer_EncDec.py new file mode 100644 index 0000000..00b3a7a --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/Autoformer_EncDec.py @@ -0,0 +1,173 @@ +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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/Embed.py b/code/new_realtime/models/third_party/patch_tst/layers/Embed.py new file mode 100644 index 0000000..abdf903 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/Embed.py @@ -0,0 +1,164 @@ +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) \ No newline at end of file diff --git a/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_backbone.py b/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_backbone.py new file mode 100644 index 0000000..ea464df --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_backbone.py @@ -0,0 +1,429 @@ +__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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_layers.py b/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_layers.py new file mode 100644 index 0000000..11b5bd6 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/PatchTST_layers.py @@ -0,0 +1,121 @@ +__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) \ No newline at end of file diff --git a/code/new_realtime/models/third_party/patch_tst/layers/RevIN.py b/code/new_realtime/models/third_party/patch_tst/layers/RevIN.py new file mode 100644 index 0000000..f7fed81 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/RevIN.py @@ -0,0 +1,63 @@ +# 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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/SelfAttention_Family.py b/code/new_realtime/models/third_party/patch_tst/layers/SelfAttention_Family.py new file mode 100644 index 0000000..c8138e2 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/SelfAttention_Family.py @@ -0,0 +1,166 @@ +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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/Transformer_EncDec.py b/code/new_realtime/models/third_party/patch_tst/layers/Transformer_EncDec.py new file mode 100644 index 0000000..c0c5789 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/layers/Transformer_EncDec.py @@ -0,0 +1,131 @@ +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 diff --git a/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_backbone.cpython-310.pyc b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_backbone.cpython-310.pyc new file mode 100644 index 0000000..5e0bf10 Binary files /dev/null and b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_backbone.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_layers.cpython-310.pyc b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_layers.cpython-310.pyc new file mode 100644 index 0000000..d287197 Binary files /dev/null and b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/PatchTST_layers.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/RevIN.cpython-310.pyc b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/RevIN.cpython-310.pyc new file mode 100644 index 0000000..784086f Binary files /dev/null and b/code/new_realtime/models/third_party/patch_tst/layers/__pycache__/RevIN.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst/models/Autoformer.py b/code/new_realtime/models/third_party/patch_tst/models/Autoformer.py new file mode 100644 index 0000000..8e66d01 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/Autoformer.py @@ -0,0 +1,121 @@ +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] diff --git a/code/new_realtime/models/third_party/patch_tst/models/DLinear.py b/code/new_realtime/models/third_party/patch_tst/models/DLinear.py new file mode 100644 index 0000000..1cf739a --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/DLinear.py @@ -0,0 +1,87 @@ +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] diff --git a/code/new_realtime/models/third_party/patch_tst/models/Informer.py b/code/new_realtime/models/third_party/patch_tst/models/Informer.py new file mode 100644 index 0000000..ef0be31 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/Informer.py @@ -0,0 +1,101 @@ +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] diff --git a/code/new_realtime/models/third_party/patch_tst/models/Linear.py b/code/new_realtime/models/third_party/patch_tst/models/Linear.py new file mode 100644 index 0000000..17095c2 --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/Linear.py @@ -0,0 +1,21 @@ +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] \ No newline at end of file diff --git a/code/new_realtime/models/third_party/patch_tst/models/NLinear.py b/code/new_realtime/models/third_party/patch_tst/models/NLinear.py new file mode 100644 index 0000000..29a0b3c --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/NLinear.py @@ -0,0 +1,24 @@ +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] \ No newline at end of file diff --git a/code/new_realtime/models/third_party/patch_tst/models/PatchTST.py b/code/new_realtime/models/third_party/patch_tst/models/PatchTST.py new file mode 100644 index 0000000..6c7f5ad --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/PatchTST.py @@ -0,0 +1,127 @@ +__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 diff --git a/code/new_realtime/models/third_party/patch_tst/models/Stat_models.py b/code/new_realtime/models/third_party/patch_tst/models/Stat_models.py new file mode 100644 index 0000000..21e853b --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/Stat_models.py @@ -0,0 +1,120 @@ +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] diff --git a/code/new_realtime/models/third_party/patch_tst/models/Transformer.py b/code/new_realtime/models/third_party/patch_tst/models/Transformer.py new file mode 100644 index 0000000..c55824d --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst/models/Transformer.py @@ -0,0 +1,94 @@ +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] diff --git a/code/new_realtime/models/third_party/patch_tst/models/__pycache__/PatchTST.cpython-310.pyc b/code/new_realtime/models/third_party/patch_tst/models/__pycache__/PatchTST.cpython-310.pyc new file mode 100644 index 0000000..3f14907 Binary files /dev/null and b/code/new_realtime/models/third_party/patch_tst/models/__pycache__/PatchTST.cpython-310.pyc differ diff --git a/code/new_realtime/models/third_party/patch_tst_raw b/code/new_realtime/models/third_party/patch_tst_raw new file mode 160000 index 0000000..204c21e --- /dev/null +++ b/code/new_realtime/models/third_party/patch_tst_raw @@ -0,0 +1 @@ +Subproject commit 204c21efe0b39603ad6e2ca640ef5896646ab1a9 diff --git a/code/new_realtime/models/third_party/tft/LICENSE b/code/new_realtime/models/third_party/tft/LICENSE new file mode 100644 index 0000000..b6dee71 --- /dev/null +++ b/code/new_realtime/models/third_party/tft/LICENSE @@ -0,0 +1,201 @@ + 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 2021-2022 NVIDIA Corporation + + 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. diff --git a/code/new_realtime/models/third_party/tft/NOTICE b/code/new_realtime/models/third_party/tft/NOTICE new file mode 100644 index 0000000..dacbdee --- /dev/null +++ b/code/new_realtime/models/third_party/tft/NOTICE @@ -0,0 +1,5 @@ +TFT for PyTorch + +This repository includes software from https://github.com/google-research/google-research/tree/master/tft licensed under the Apache 2.0 License. + +This repository contains code from https://github.com/rwightman/pytorch-image-models/blob/master/timm/utils/model_ema.py under the Apache 2.0 License. diff --git a/code/new_realtime/models/third_party/tft/README.md b/code/new_realtime/models/third_party/tft/README.md new file mode 100644 index 0000000..65982eb --- /dev/null +++ b/code/new_realtime/models/third_party/tft/README.md @@ -0,0 +1,3 @@ +This folder contains code copied from NVIDIA's Temporal Fusion Transformer implementation, licensed under Apache 2.0. +All rights belong to NVIDIA Corporation. +Modifications are noted in the file headers. \ No newline at end of file diff --git a/code/new_realtime/models/third_party/tft/__init__.py b/code/new_realtime/models/third_party/tft/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/code/new_realtime/models/third_party/tft/tft_model.py b/code/new_realtime/models/third_party/tft/tft_model.py new file mode 100755 index 0000000..99b85df --- /dev/null +++ b/code/new_realtime/models/third_party/tft/tft_model.py @@ -0,0 +1,525 @@ +# Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +# Modified by Alexander Blank, 2025. +# Modifications: +# - added support for multiple outputs +# - added support for mode configurable targets +# - added support for single dimension, non-quantile outputs +# - added support for target agnostic predictions, for cases, where the target does not become known after prediction + + +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torch import Tensor +from torch.nn.parameter import UninitializedParameter +from typing import Dict, Tuple, Optional, List + + +MAKE_CONVERT_COMPATIBLE = os.environ.get("TFT_SCRIPTING", None) is not None +from torch.nn import LayerNorm + + +class MaybeLayerNorm(nn.Module): + def __init__(self, output_size, hidden_size, eps): + super().__init__() + if output_size and output_size == 1: + self.ln = nn.Identity() + else: + self.ln = LayerNorm(output_size if output_size else hidden_size, eps=eps) + + def forward(self, x): + return self.ln(x) + + +class GLU(nn.Module): + def __init__(self, hidden_size, output_size): + super().__init__() + self.lin = nn.Linear(hidden_size, output_size * 2) + + def forward(self, x: Tensor) -> Tensor: + x = self.lin(x) + x = F.glu(x) + return x + + +class GRN(nn.Module): + def __init__(self, + input_size, + hidden_size, + output_size=None, + context_hidden_size=None, + dropout=0.0, ): + super().__init__() + self.layer_norm = MaybeLayerNorm(output_size, hidden_size, eps=1e-3) + self.lin_a = nn.Linear(input_size, hidden_size) + if context_hidden_size is not None: + self.lin_c = nn.Linear(context_hidden_size, hidden_size, bias=False) + else: + self.lin_c = nn.Identity() + self.lin_i = nn.Linear(hidden_size, hidden_size) + self.glu = GLU(hidden_size, output_size if output_size else hidden_size) + self.dropout = nn.Dropout(dropout) + self.out_proj = nn.Linear(input_size, output_size) if output_size else None + + def forward(self, a: Tensor, c: Optional[Tensor] = None): + x = self.lin_a(a) + if c is not None: + x = x + self.lin_c(c).unsqueeze(1) + x = F.elu(x) + x = self.lin_i(x) + x = self.dropout(x) + x = self.glu(x) + y = a if self.out_proj is None else self.out_proj(a) + x = x + y + return self.layer_norm(x) + + # @torch.jit.script #Currently broken with autocast + + +def fused_pointwise_linear_v1(x, a, b): + out = torch.mul(x.unsqueeze(-1), a) + out = out + b + return out + + +@torch.jit.script +def fused_pointwise_linear_v2(x, a, b): + out = x.unsqueeze(3) * a + out = out + b + return out + + +class TFTEmbedding(nn.Module): + def __init__(self, config, initialize_cont_params=True): + # initialize_cont_params=False prevents form initializing parameters inside this class + # so they can be lazily initialized in LazyEmbedding module + super().__init__() + self.s_cat_inp_lens = config.static_categorical_inp_lens + self.t_cat_k_inp_lens = config.temporal_known_categorical_inp_lens + self.t_cat_o_inp_lens = config.temporal_observed_categorical_inp_lens + self.s_cont_inp_size = config.static_continuous_inp_size + self.t_cont_k_inp_size = config.temporal_known_continuous_inp_size + self.t_cont_o_inp_size = config.temporal_observed_continuous_inp_size + self.t_tgt_size = config.temporal_target_size + + self.hidden_size = config.hidden_size + + # There are 7 types of input: + # 1. Static categorical + # 2. Static continuous + # 3. Temporal known a priori categorical + # 4. Temporal known a priori continuous + # 5. Temporal observed categorical + # 6. Temporal observed continuous + # 7. Temporal observed targets (time series obseved so far) + + self.s_cat_embed = nn.ModuleList([ + nn.Embedding(n, self.hidden_size) for n in self.s_cat_inp_lens]) if self.s_cat_inp_lens else None + self.t_cat_k_embed = nn.ModuleList([ + nn.Embedding(n, self.hidden_size) for n in self.t_cat_k_inp_lens]) if self.t_cat_k_inp_lens else None + self.t_cat_o_embed = nn.ModuleList([ + nn.Embedding(n, self.hidden_size) for n in self.t_cat_o_inp_lens]) if self.t_cat_o_inp_lens else None + + if initialize_cont_params: + self.s_cont_embedding_vectors = nn.Parameter( + torch.Tensor(self.s_cont_inp_size, self.hidden_size)) if self.s_cont_inp_size else None + self.t_cont_k_embedding_vectors = nn.Parameter( + torch.Tensor(self.t_cont_k_inp_size, self.hidden_size)) if self.t_cont_k_inp_size else None + self.t_cont_o_embedding_vectors = nn.Parameter( + torch.Tensor(self.t_cont_o_inp_size, self.hidden_size)) if self.t_cont_o_inp_size else None + self.t_tgt_embedding_vectors = nn.Parameter(torch.Tensor(self.t_tgt_size, self.hidden_size)) + + self.s_cont_embedding_bias = nn.Parameter( + torch.zeros(self.s_cont_inp_size, self.hidden_size)) if self.s_cont_inp_size else None + self.t_cont_k_embedding_bias = nn.Parameter( + torch.zeros(self.t_cont_k_inp_size, self.hidden_size)) if self.t_cont_k_inp_size else None + self.t_cont_o_embedding_bias = nn.Parameter( + torch.zeros(self.t_cont_o_inp_size, self.hidden_size)) if self.t_cont_o_inp_size else None + self.t_tgt_embedding_bias = nn.Parameter(torch.zeros(self.t_tgt_size, self.hidden_size)) + + self.reset_parameters() + + def reset_parameters(self): + if self.s_cont_embedding_vectors is not None: + torch.nn.init.xavier_normal_(self.s_cont_embedding_vectors) + torch.nn.init.zeros_(self.s_cont_embedding_bias) + if self.t_cont_k_embedding_vectors is not None: + torch.nn.init.xavier_normal_(self.t_cont_k_embedding_vectors) + torch.nn.init.zeros_(self.t_cont_k_embedding_bias) + if self.t_cont_o_embedding_vectors is not None: + torch.nn.init.xavier_normal_(self.t_cont_o_embedding_vectors) + torch.nn.init.zeros_(self.t_cont_o_embedding_bias) + if self.t_tgt_embedding_vectors is not None: + torch.nn.init.xavier_normal_(self.t_tgt_embedding_vectors) + torch.nn.init.zeros_(self.t_tgt_embedding_bias) + if self.s_cat_embed is not None: + for module in self.s_cat_embed: + module.reset_parameters() + if self.t_cat_k_embed is not None: + for module in self.t_cat_k_embed: + module.reset_parameters() + if self.t_cat_o_embed is not None: + for module in self.t_cat_o_embed: + module.reset_parameters() + + def _apply_embedding(self, + cat: Optional[Tensor], + cont: Optional[Tensor], + cat_emb: Optional[nn.ModuleList], + cont_emb: Tensor, + cont_bias: Tensor, + ) -> Tuple[Optional[Tensor], Optional[Tensor]]: + e_cat = torch.stack([embed(cat[..., i]) for i, embed in enumerate(cat_emb)], + dim=-2) if cat is not None else None + if cont is not None: + # the line below is equivalent to following einsums + # e_cont = torch.einsum('btf,fh->bthf', cont, cont_emb) + # e_cont = torch.einsum('bf,fh->bhf', cont, cont_emb) + if MAKE_CONVERT_COMPATIBLE: + e_cont = torch.mul(cont.unsqueeze(-1), cont_emb) + e_cont = e_cont + cont_bias + else: + e_cont = fused_pointwise_linear_v1(cont, cont_emb, cont_bias) + else: + e_cont = None + + if e_cat is not None and e_cont is not None: + return torch.cat([e_cat, e_cont], dim=-2) + elif e_cat is not None: + return e_cat + elif e_cont is not None: + return e_cont + else: + return None + + def forward(self, x: Dict[str, Tensor], use_target: bool = False): + # Extract inputs + s_cat_inp = x.get('s_cat', None) + s_cont_inp = x.get('s_cont', None) + t_cat_k_inp = x.get('k_cat', None) + t_cont_k_inp = x.get('k_cont', None) + t_cat_o_inp = x.get('o_cat', None) + t_cont_o_inp = x.get('o_cont', None) + + # Only use target if teacher forcing is enabled. + # When disabled, we ignore target values. + if use_target: + t_tgt_obs = x['target'] # Must be present when using teacher forcing + else: + t_tgt_obs = None + + # For static inputs, take the first timestep + s_cat_inp = s_cat_inp[:, 0, :] if s_cat_inp is not None else None + s_cont_inp = s_cont_inp[:, 0, :] if s_cont_inp is not None else None + + # Apply embeddings for static and known/observed temporal features + s_inp = self._apply_embedding(s_cat_inp, + s_cont_inp, + self.s_cat_embed, + self.s_cont_embedding_vectors, + self.s_cont_embedding_bias) + t_known_inp = self._apply_embedding(t_cat_k_inp, + t_cont_k_inp, + self.t_cat_k_embed, + self.t_cont_k_embedding_vectors, + self.t_cont_k_embedding_bias) + t_observed_inp = self._apply_embedding(t_cat_o_inp, + t_cont_o_inp, + self.t_cat_o_embed, + self.t_cont_o_embedding_vectors, + self.t_cont_o_embedding_bias) + # Compute the target embedding only if teacher forcing is enabled. + if use_target and t_tgt_obs is not None: + if MAKE_CONVERT_COMPATIBLE: + t_observed_tgt = torch.matmul(t_tgt_obs.unsqueeze(3).unsqueeze(4), + self.t_tgt_embedding_vectors.unsqueeze(1)).squeeze(3) + t_observed_tgt = t_observed_tgt + self.t_tgt_embedding_bias + else: + t_observed_tgt = fused_pointwise_linear_v2(t_tgt_obs, + self.t_tgt_embedding_vectors, + self.t_tgt_embedding_bias) + else: + t_observed_tgt = None + + return s_inp, t_known_inp, t_observed_inp, t_observed_tgt + + +class LazyEmbedding(nn.modules.lazy.LazyModuleMixin, TFTEmbedding): + cls_to_become = TFTEmbedding + + def __init__(self, config): + super().__init__(config, initialize_cont_params=False) + + if config.static_continuous_inp_size: + self.s_cont_embedding_vectors = UninitializedParameter() + self.s_cont_embedding_bias = UninitializedParameter() + else: + self.s_cont_embedding_vectors = None + self.s_cont_embedding_bias = None + + if config.temporal_known_continuous_inp_size: + self.t_cont_k_embedding_vectors = UninitializedParameter() + self.t_cont_k_embedding_bias = UninitializedParameter() + else: + self.t_cont_k_embedding_vectors = None + self.t_cont_k_embedding_bias = None + + if config.temporal_observed_continuous_inp_size: + self.t_cont_o_embedding_vectors = UninitializedParameter() + self.t_cont_o_embedding_bias = UninitializedParameter() + else: + self.t_cont_o_embedding_vectors = None + self.t_cont_o_embedding_bias = None + + self.t_tgt_embedding_vectors = UninitializedParameter() + self.t_tgt_embedding_bias = UninitializedParameter() + + def initialize_parameters(self, x): + if self.has_uninitialized_params(): + s_cont_inp = x.get('s_cont', None) + t_cont_k_inp = x.get('k_cont', None) + t_cont_o_inp = x.get('o_cont', None) + t_tgt_obs = x['target'] # Has to be present + + if s_cont_inp is not None: + self.s_cont_embedding_vectors.materialize((s_cont_inp.shape[-1], self.hidden_size)) + self.s_cont_embedding_bias.materialize((s_cont_inp.shape[-1], self.hidden_size)) + + if t_cont_k_inp is not None: + self.t_cont_k_embedding_vectors.materialize((t_cont_k_inp.shape[-1], self.hidden_size)) + self.t_cont_k_embedding_bias.materialize((t_cont_k_inp.shape[-1], self.hidden_size)) + + if t_cont_o_inp is not None: + self.t_cont_o_embedding_vectors.materialize((t_cont_o_inp.shape[-1], self.hidden_size)) + self.t_cont_o_embedding_bias.materialize((t_cont_o_inp.shape[-1], self.hidden_size)) + + self.t_tgt_embedding_vectors.materialize((t_tgt_obs.shape[-1], self.hidden_size)) + self.t_tgt_embedding_bias.materialize((t_tgt_obs.shape[-1], self.hidden_size)) + + self.reset_parameters() + + # def forward(self, x: Dict[str, Tensor], use_target: bool = True): + # return super().forward(x, use_target=use_target) + + +class VariableSelectionNetwork(nn.Module): + def __init__(self, config, num_inputs): + super().__init__() + self.joint_grn = GRN(config.hidden_size * num_inputs, config.hidden_size, output_size=num_inputs, + context_hidden_size=config.hidden_size) + self.var_grns = nn.ModuleList( + [GRN(config.hidden_size, config.hidden_size, dropout=config.dropout) for _ in range(num_inputs)]) + + def forward(self, x: Tensor, context: Optional[Tensor] = None): + Xi = torch.flatten(x, start_dim=-2) + grn_outputs = self.joint_grn(Xi, c=context) + sparse_weights = F.softmax(grn_outputs, dim=-1) + transformed_embed_list = [m(x[..., i, :]) for i, m in enumerate(self.var_grns)] + transformed_embed = torch.stack(transformed_embed_list, dim=-1) + # the line below performs batched matrix vector multiplication + # for temporal features it's bthf,btf->bth + # for static features it's bhf,bf->bh + variable_ctx = torch.matmul(transformed_embed, sparse_weights.unsqueeze(-1)).squeeze(-1) + + return variable_ctx, sparse_weights + + +class StaticCovariateEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.vsn = VariableSelectionNetwork(config, config.num_static_vars) + self.context_grns = nn.ModuleList( + [GRN(config.hidden_size, config.hidden_size, dropout=config.dropout) for _ in range(4)]) + + def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + variable_ctx, sparse_weights = self.vsn(x) + + # Context vectors: + # variable selection context + # enrichment context + # state_c context + # state_h context + cs, ce, ch, cc = [m(variable_ctx) for m in self.context_grns] + + return cs, ce, ch, cc + + +class InterpretableMultiHeadAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.n_head = config.n_head + assert config.hidden_size % config.n_head == 0 + self.d_head = config.hidden_size // config.n_head + self.qkv_linears = nn.Linear(config.hidden_size, (2 * self.n_head + 1) * self.d_head, bias=False) + self.out_proj = nn.Linear(self.d_head, config.hidden_size, bias=False) + self.attn_dropout = nn.Dropout(config.attn_dropout) + self.out_dropout = nn.Dropout(config.dropout) + self.scale = self.d_head ** -0.5 + self.register_buffer("_mask", + torch.triu(torch.full((config.example_length, config.example_length), float('-inf')), + 1).unsqueeze(0)) + + def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]: + bs, t, h_size = x.shape + qkv = self.qkv_linears(x) + q, k, v = qkv.split((self.n_head * self.d_head, self.n_head * self.d_head, self.d_head), dim=-1) + q = q.view(bs, t, self.n_head, self.d_head) + k = k.view(bs, t, self.n_head, self.d_head) + v = v.view(bs, t, self.d_head) + + # attn_score = torch.einsum('bind,bjnd->bnij', q, k) + attn_score = torch.matmul(q.permute((0, 2, 1, 3)), k.permute((0, 2, 3, 1))) + attn_score.mul_(self.scale) + + attn_score = attn_score + self._mask + + attn_prob = F.softmax(attn_score, dim=3) + attn_prob = self.attn_dropout(attn_prob) + + # attn_vec = torch.einsum('bnij,bjd->bnid', attn_prob, v) + attn_vec = torch.matmul(attn_prob, v.unsqueeze(1)) + m_attn_vec = torch.mean(attn_vec, dim=1) + out = self.out_proj(m_attn_vec) + out = self.out_dropout(out) + + return out, attn_prob + + +class TFTBack(nn.Module): + def __init__(self, config): + super().__init__() + + self.encoder_length = config.encoder_length + self.history_vsn = VariableSelectionNetwork(config, config.num_historic_vars) + self.history_encoder = nn.LSTM(config.hidden_size, config.hidden_size, batch_first=True) + self.future_vsn = VariableSelectionNetwork(config, config.num_future_vars) + self.future_encoder = nn.LSTM(config.hidden_size, config.hidden_size, batch_first=True) + + self.input_gate = GLU(config.hidden_size, config.hidden_size) + self.input_gate_ln = LayerNorm(config.hidden_size, eps=1e-3) + + self.enrichment_grn = GRN(config.hidden_size, + config.hidden_size, + context_hidden_size=config.hidden_size, + dropout=config.dropout) + self.attention = InterpretableMultiHeadAttention(config) + self.attention_gate = GLU(config.hidden_size, config.hidden_size) + self.attention_ln = LayerNorm(config.hidden_size, eps=1e-3) + + self.positionwise_grn = GRN(config.hidden_size, + config.hidden_size, + dropout=config.dropout) + + self.decoder_gate = GLU(config.hidden_size, config.hidden_size) + self.decoder_ln = LayerNorm(config.hidden_size, eps=1e-3) + + self.quantiles = config.quantiles + self.target_size = config.target_size + if self.quantiles is not None: + self.output = nn.Linear(config.hidden_size, len(config.quantiles) * config.target_size) + else: + self.output = nn.Linear(config.hidden_size, config.target_size) + + def forward(self, historical_inputs, cs, ch, cc, ce, future_inputs): + historical_features, _ = self.history_vsn(historical_inputs, cs) + history, state = self.history_encoder(historical_features, (ch, cc)) + future_features, _ = self.future_vsn(future_inputs, cs) + future, _ = self.future_encoder(future_features, state) + torch.cuda.synchronize() + + # skip connection + input_embedding = torch.cat([historical_features, future_features], dim=1) + temporal_features = torch.cat([history, future], dim=1) + temporal_features = self.input_gate(temporal_features) + temporal_features = temporal_features + input_embedding + temporal_features = self.input_gate_ln(temporal_features) + + # Static enrichment + enriched = self.enrichment_grn(temporal_features, c=ce) + + # Temporal self attention + x, _ = self.attention(enriched) + + # Don't compute hictorical quantiles + x = x[:, self.encoder_length:, :] + temporal_features = temporal_features[:, self.encoder_length:, :] + enriched = enriched[:, self.encoder_length:, :] + + x = self.attention_gate(x) + x = x + enriched + x = self.attention_ln(x) + + # Position-wise feed-forward + x = self.positionwise_grn(x) + + # Final skip connection + x = self.decoder_gate(x) + x = x + temporal_features + x = self.decoder_ln(x) + + out = self.output(x) + if self.quantiles is not None: + # Reshape to [batch, time, target_size, n_quantiles] + out = out.view(out.size(0), out.size(1), self.target_size, len(self.quantiles)) + else: + # Reshape to [batch, time, target_size] + out = out.view(out.size(0), out.size(1), self.target_size) + + return out + + +class TemporalFusionTransformer(nn.Module): + """ + Implementation of https://arxiv.org/abs/1912.09363 + """ + + def __init__(self, config): + super().__init__() + + if hasattr(config, 'model'): + config = config.model + + self.encoder_length = config.encoder_length # this determines from how distant past we want to use data from + + # self.embedding = LazyEmbedding(config) + self.embedding = TFTEmbedding(config) + self.static_encoder = StaticCovariateEncoder(config) + # if MAKE_CONVERT_COMPATIBLE: + self.TFTpart2 = TFTBack(config) + # else: + # self.TFTpart2 = torch.jit.script(TFTBack(config)) + + def forward(self, x: Dict[str, Tensor]) -> Tensor: + # Call embedding with use_target=False to skip target features entirely. + s_inp, t_known_inp, t_observed_inp, t_observed_tgt = self.embedding(x, use_target=False) + + # Compute static context + cs, ce, ch, cc = self.static_encoder(s_inp) + ch, cc = ch.unsqueeze(0), cc.unsqueeze(0) # Initialize LSTM states + + # Build historical inputs without teacher-forced targets. + # Include observed features if available, and the known inputs. + historical_inputs = [] + if t_observed_inp is not None: + historical_inputs.append(t_observed_inp[:, :self.encoder_length, :]) + historical_inputs.append(t_known_inp[:, :self.encoder_length, :]) + historical_inputs = torch.cat(historical_inputs, dim=-2) + + # Future inputs remain the same + future_inputs = t_known_inp[:, self.encoder_length:] + return self.TFTpart2(historical_inputs, cs, ch, cc, ce, future_inputs) diff --git a/code/new_realtime/models/transformer.py b/code/new_realtime/models/transformer.py new file mode 100644 index 0000000..9ec0df7 --- /dev/null +++ b/code/new_realtime/models/transformer.py @@ -0,0 +1,40 @@ +import math + +import torch +from torch import nn + + +class TransformerModel(nn.Module): + def __init__(self, input_dim: int, + output_dim: int, + seq_len: int, + embed_dim: int, + num_heads: int, + num_enc_layers: int): + super().__init__() + 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) + self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers) + self.pool = nn.AdaptiveAvgPool1d(1) + self.head = nn.Linear(embed_dim, output_dim) + + def _get_sinusoidal_embedding(self, seq_len, embed_dim): + position = torch.arange(0, seq_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim)) + pe = torch.zeros(seq_len, embed_dim) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + return pe # (seq_len, embed_dim) + + def forward(self, x): + # x: (B, seq_len, input_dim) + x = self.input_proj(x) + self.pos_embed[:, :x.size(1), :] # broadcasting + x = x.permute(1, 0, 2) # (S, B, E) + enc = self.encoder(x) # (S, B, E) + pooled = enc.mean(0) # (B, E) + return self.head(pooled) diff --git a/code/new_realtime/models/utils.py b/code/new_realtime/models/utils.py new file mode 100644 index 0000000..a01de0a --- /dev/null +++ b/code/new_realtime/models/utils.py @@ -0,0 +1,137 @@ +import os +from typing import Callable +import math + +import numpy as np +from bson import ObjectId +import torch +from torch import nn + +from utils.data_utils import get_collated_batch_for_key + + +def simple_get_y(collated_batch): + """ + Get y from collated batch + Args: + collated_batch: collated batch to get y from + + Returns: + np.ndarray: y + """ + return collated_batch[1] + + +def simple_model_save(model: nn.Module, + training_configuration: dict) -> None: + """ + Save model to disk + Args: + model: model to save + training_configuration: training configuration + Returns: + None + """ + model_path = os.path.join(training_configuration["training_dir"], "model.pt") + torch.save(model.state_dict(), model_path) + + +def simple_model_load(model_configuration: dict, + training_configuration: dict, + sample_key: str | ObjectId, + device: str, + model_creation_fn: Callable, + *args, **kwargs) -> nn.Module: + """ + Load model from disk + Args: + model_configuration: model configuration + training_configuration: training configuration + sample_key: sample key to get sample data batch with + device: device to load model on + model_creation_fn: function to create model + Returns: + model: loaded model + """ + training_dir = training_configuration["training_dir"] + model_state_path = os.path.join(training_dir, "model.pt") + model = model_creation_fn(model_configuration, + sample_key, *args, **kwargs) + model.to(device) + model.load_state_dict(torch.load(model_state_path)) + return model + + +def simple_model_creation(model_configuration: dict, + sample_key: str | ObjectId, + lmdb_env=None) -> nn.Module: + """ + Create model from configuration + Args: + model_configuration: model configuration + sample_key: sample key to get sample data batch with + lmdb_env: LMDB environment to use for getting sample data batch + + Returns: + model: created model + """ + sample_item = get_collated_batch_for_key(sample_key, model_configuration, lmdb_env=lmdb_env) + input_size = sample_item[0].shape[2] + output_size = sample_item[1].shape[2] + model_class = model_configuration["model_class"] + cnn_model = model_class(input_dim=input_size, + output_dim=output_size, + **model_configuration["model_parameters"], ) + return cnn_model + + +def simple_x_y_predict(model: nn.Module, + collated_batch: tuple[torch.Tensor, torch.Tensor], + device: str, + batch_size: int, + *args, **kwargs) -> torch.Tensor: + """ + Predict y from x + Args: + model: model to use for prediction, must have 'predict' method + collated_batch: collated data batch to predict with, + must be a tuple of (x, y) where x is the input data and y is the target data + device: device to use for prediction + *args: additional arguments to pass to the model's predict method + **kwargs: additional keyword arguments to pass to the model's predict method + + Returns: + torch.Tensor: predicted y as numpy array on CPU + """ + + model.eval() + x = collated_batch[0].to(device).float() + batches = list() + num_batches = math.ceil(len(x) / batch_size) + for i in range(num_batches): + start = i * batch_size + end = (i + 1) * batch_size + if end > len(x): + end = len(x) + batch_slice = x[start:end] + batches.append(batch_slice) + + preds = list() + with torch.no_grad(): + for batch in batches: + batch = batch.to(device).float() + pred = model(batch, *args, **kwargs) + preds.append(pred.cpu().numpy()) + return np.concatenate(preds) + + +def simple_get_y(collated_batch: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: + """ + Get y from collated batch + Args: + collated_batch: collated batch to get y from + + Returns: + torch.Tensor: y + """ + return collated_batch[1] diff --git a/code/new_realtime/training_wrapper.py b/code/new_realtime/training_wrapper.py new file mode 100644 index 0000000..ac7b5b5 --- /dev/null +++ b/code/new_realtime/training_wrapper.py @@ -0,0 +1,321 @@ +import json +import sys +import os +import argparse +import logging +from datetime import datetime + +import lmdb +import dotenv +import torch +from torch.utils.data import DataLoader +import torch.distributed as dist + +from experiment_setup import get_eval_functions +from utils.evaluation import evaluate_model +from utils.model_utils import get_model_config +from utils.training_utils import get_training_config, get_data_ids +from utils.data_utils import LMDBIterableDataset +from utils.utils import get_variable_from_module, get_logger, convert_for_json +from utils.training import train_model + +dotenv.load_dotenv() + +logger = None + +# set up distributed training +dist.init_process_group(backend="nccl", init_method="env://") +local_rank = torch.distributed.get_rank() +torch.cuda.set_device(local_rank) + + +def prepare_run(results_dir: str, + lmdb_root_dir: str, + base_model_configuration: dict, + base_training_configuration: dict): + # create model configuration from base + logger.info("Creating model configuration") + model_configuration = get_model_config(base_model_configuration, results_dir, lmdb_root_dir) + feature_config = model_configuration["feature_config"] + dataset_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}" + logger.info(f"Model configuration name: {model_configuration['id']}") + + # create training configuration + logger.info("Creating training configuration") + training_configuration = get_training_config(base_training_configuration, model_configuration) + logger.info(f"Training configuration: {training_configuration['id']}") + + logger.info(f"Fetching data ids, limit: {item_limit}") + train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, item_limit) + + logger.info(f"Train ids: {len(train_ids)}, Val ids: {len(val_ids)}, Test ids: {len(test_ids)}") + + return model_configuration, training_configuration, train_ids, val_ids, test_ids, dataset_dir + + +def train( + model_configuration: dict, + training_configuration: dict, + train_ids: list, + val_ids: list, + dataset_dir: str, + log_dir: str) -> None: + # create training and validation loaders + logger.info("Creating training loaders") + train_loader = LMDBIterableDataset(dataset_dir, + train_ids, + model_configuration=model_configuration, + batch_size=training_configuration["batch_size"]) + logger.info("Creating validation loaders") + val_loader = LMDBIterableDataset(dataset_dir, + val_ids, + model_configuration=model_configuration, + batch_size=training_configuration["batch_size"]) + + # instantiate model + logger.info("Creating model") + lmdb_env = lmdb.open(dataset_dir, readonly=True) + model_creation_fn = model_configuration["model_creation_fn"] + model = model_creation_fn(model_configuration, + train_ids[0], + lmdb_env=lmdb_env) + + logger.info(f"Started training for model {model_configuration['id']}") + train_model( + model=model, + model_configuration=model_configuration, + training_configuration=training_configuration, + train_dataset=train_loader, + val_dataset=val_loader, + log_dir=log_dir, + logger=logger + ) + + +def evaluate(model_configuration: dict, + training_configuration: dict, + test_ids: list, + device) -> dict: + logger.info(f"Evaluating model {model_configuration['id']}") + eval_functions = get_eval_functions(model_configuration) + results = evaluate_model(model_configuration, + training_configuration, + test_ids, + eval_functions) + + # save results to file + results_dir = training_configuration["training_dir"] + if not os.path.exists(results_dir): + os.makedirs(results_dir) + results_file = os.path.join(results_dir, "evaluation_results.json") + with open(results_file, "w") as f: + json.dump(convert_for_json(results), f, indent=4) + return results + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Training wrapper for model training") + parser.add_argument("run_configuration_module", + type=str, + help="Path to the run configuration module") + parser.add_argument("--run_configuration_variable", + type=str, + required=False, + default="run_configuration", + help="Name of the run configuration variable in the module") + parser.add_argument("--results_dir", + type=str, + required=False, + default=None, + help="Directory to save the results") + parser.add_argument("--lmdb_root_dir", + type=str, + required=False, + default=None, + help="Path to the lmdb directory") + parser.add_argument("--log_dir", + type=str, + required=False, + default=None, + help="Directory to save the logs") + + parser.add_argument("--item_limit", + type=int, + required=False, + default=None, + help="Limit the number of items to process, default is None (no limit)") + parser.add_argument("--device", + type=str, + required=False, + default="cuda", + help="Device to use for training, default is cuda") + + args = parser.parse_args() + # load run configuration from module + run_configuration = get_variable_from_module( + # make sure to replace / with . and remove .py to get proper module tree + module_path=args.run_configuration_module.replace(".py", "").replace("/", "."), + variable_name=args.run_configuration_variable) + + # check for variables in run_configuration + if not args.results_dir: + if "base_results_dir" not in run_configuration: + results_dir = os.getenv("RESULTS_ROOT_DIR") + else: + results_dir = run_configuration["base_results_dir"] + else: + results_dir = args.results_dir + + if results_dir is None: + raise ValueError( + "No results directory specified. Please set the RESULTS_ROOT_DIR environment variable or provide a results_dir argument.") + + if not os.path.exists(results_dir): + os.makedirs(results_dir) + + if not args.lmdb_root_dir: + if "base_lmdb_root_dir" not in run_configuration: + lmdb_root_dir = os.getenv("LMDB_ROOT_DIR") + else: + lmdb_root_dir = run_configuration["base_lmdb_root_dir"] + else: + lmdb_root_dir = args.lmdb_root_dir + + if lmdb_root_dir is None: + raise ValueError( + "No LMDB root directory specified. Please set the LMDB_ROOT_DIR environment variable or provide a lmdb_root_dir argument.") + + if not os.path.exists(lmdb_root_dir): + os.makedirs(lmdb_root_dir) + + if not args.log_dir: + if "base_log_dir" not in run_configuration: + log_dir = os.getenv("LOG_DIR") + else: + log_dir = run_configuration["log_dir"] + else: + log_dir = args.log_dir + if log_dir is None: + raise ValueError( + "No log directory specified. Please set the LOG_DIR environment variable or provide a log_dir argument.") + if not os.path.exists(log_dir): + os.makedirs(log_dir) + + item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"] + if item_limit == -1: + item_limit = None + + run_name = run_configuration["name"] + runs = run_configuration["runs"] + + # set up run_id on rank 0 + if local_rank == 0: + run_id = f"{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + else: + run_id = None + + # broadcast run_id to all processes + if dist.is_initialized(): + run_id_list = [run_id] + torch.distributed.broadcast_object_list(run_id_list, src=0) + run_id = run_id_list[0] + + # make sure run_id is a string + run_id = str(run_id) + + + # append run_id to results_dir and log_dir + results_dir = os.path.join(results_dir, run_id) + log_dir = os.path.join(log_dir, run_id) + + # create directories if they do not exist, only on the main process + if torch.distributed.get_rank() == 0 or not dist.is_initialized(): + if not os.path.exists(results_dir): + os.makedirs(results_dir) + if not os.path.exists(log_dir): + os.makedirs(log_dir) + + # set up logger + logger = get_logger(module_name=run_name, filename=os.path.join(log_dir, "main.log")) + + logger.info(f"Rank {local_rank}: Starting run {run_name}") + logger.info(f"Rank {local_rank}: Run ID: {run_id}") + + # print parameter values + logger.info(f"Rank {local_rank}: Results directory: {results_dir}") + logger.info(f"Rank {local_rank}: LMDB root directory: {lmdb_root_dir}") + logger.info(f"Rank {local_rank}: Log directory: {log_dir}") + logger.info(f"Rank {local_rank}: Item limit: {item_limit}") + + if not runs or len(runs) == 0: + raise ValueError("No runs specified in the run configuration. Please provide a list of runs to train.") + + try: + for run in runs: + run_step_name = run["name"] + run_description = run["description"] + run_model_configuration = run["model_configuration"] + run_training_configuration = run["training_configuration"] + + logger.info(f"Rank {local_rank}: Running: {run_step_name}") + logger.info(f"Rank {local_rank}: Description: {run_description}") + + # prepare run, make sure rank 0 is the first to avoid race conditions + if local_rank == 0: + logger.info(f"Rank {local_rank}: Preparing run {run_step_name}") + run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run( + results_dir=results_dir, + lmdb_root_dir=lmdb_root_dir, + base_model_configuration=run_model_configuration, + base_training_configuration=run_training_configuration + ) + logger.info(f"Rank {local_rank}: Finished preparing run {run_step_name}") + # sync after preparing run + if dist.is_initialized(): + dist.barrier() + else: + # wait for rank 0 to finish preparing run + if dist.is_initialized(): + dist.barrier() + logger.info(f"Rank {local_rank}: Waiting for rank 0 to finish preparing run {run_step_name}") + run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run( + results_dir=results_dir, + lmdb_root_dir=lmdb_root_dir, + base_model_configuration=run_model_configuration, + base_training_configuration=run_training_configuration + ) + + train( + model_configuration=run_model_configuration, + training_configuration=run_training_configuration, + train_ids=train_ids, + val_ids=val_ids, + dataset_dir=dataset_dir, + log_dir=log_dir, + ) + # sync after training + if dist.is_initialized(): + dist.barrier() + + logger.info(f"Rank {torch.distributed.get_rank()} finished training {run_step_name}") + + # run evaluation, only on rank 0 + local_rank = torch.distributed.get_rank() + if local_rank == 0: + logger.info(f"Rank {local_rank} starting evaluation for {run_step_name}") + results = evaluate( + model_configuration=run_model_configuration, + training_configuration=run_training_configuration, + test_ids=test_ids, + ) + logger.info( + f"Rank {local_rank} finished evaluation for {run_step_name} with model {run_model_configuration['id']}") + logger.info(f"Results: {results}") + + # sync after evaluation + if dist.is_initialized(): + dist.barrier() + finally: + # clean up + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/code/new_realtime/utils/__init__.py b/code/new_realtime/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/code/new_realtime/utils/__pycache__/__init__.cpython-310.pyc b/code/new_realtime/utils/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..7e8a5d8 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/__init__.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/data_utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/data_utils.cpython-310.pyc new file mode 100644 index 0000000..142c8c8 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/data_utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/dataset_creation.cpython-310.pyc b/code/new_realtime/utils/__pycache__/dataset_creation.cpython-310.pyc new file mode 100644 index 0000000..5400ba6 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/dataset_creation.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/dataset_utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/dataset_utils.cpython-310.pyc new file mode 100644 index 0000000..00f2490 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/dataset_utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/evaluation.cpython-310.pyc b/code/new_realtime/utils/__pycache__/evaluation.cpython-310.pyc new file mode 100644 index 0000000..1d83134 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/evaluation.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/feature_functions.cpython-310.pyc b/code/new_realtime/utils/__pycache__/feature_functions.cpython-310.pyc new file mode 100644 index 0000000..301b7e9 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/feature_functions.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/lmdb_utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/lmdb_utils.cpython-310.pyc new file mode 100644 index 0000000..1f9edd3 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/lmdb_utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/loss_functions.cpython-310.pyc b/code/new_realtime/utils/__pycache__/loss_functions.cpython-310.pyc new file mode 100644 index 0000000..4d08ced Binary files /dev/null and b/code/new_realtime/utils/__pycache__/loss_functions.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/model_utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/model_utils.cpython-310.pyc new file mode 100644 index 0000000..6926928 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/model_utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/smoothing.cpython-310.pyc b/code/new_realtime/utils/__pycache__/smoothing.cpython-310.pyc new file mode 100644 index 0000000..cc97e73 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/smoothing.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/training.cpython-310.pyc b/code/new_realtime/utils/__pycache__/training.cpython-310.pyc new file mode 100644 index 0000000..6a9807c Binary files /dev/null and b/code/new_realtime/utils/__pycache__/training.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/training_utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/training_utils.cpython-310.pyc new file mode 100644 index 0000000..6c4b92f Binary files /dev/null and b/code/new_realtime/utils/__pycache__/training_utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/utils.cpython-310.pyc b/code/new_realtime/utils/__pycache__/utils.cpython-310.pyc new file mode 100644 index 0000000..c724c82 Binary files /dev/null and b/code/new_realtime/utils/__pycache__/utils.cpython-310.pyc differ diff --git a/code/new_realtime/utils/__pycache__/visualization.cpython-310.pyc b/code/new_realtime/utils/__pycache__/visualization.cpython-310.pyc new file mode 100644 index 0000000..598c76e Binary files /dev/null and b/code/new_realtime/utils/__pycache__/visualization.cpython-310.pyc differ diff --git a/code/new_realtime/utils/data_utils.py b/code/new_realtime/utils/data_utils.py new file mode 100644 index 0000000..e0bba83 --- /dev/null +++ b/code/new_realtime/utils/data_utils.py @@ -0,0 +1,576 @@ +import copy +import logging +import os +import random + +import lmdb +import numpy as np +from bson import ObjectId +from torch.utils.data import IterableDataset +from tqdm import tqdm + +from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features +from utils.dataset_utils import load_key_stats +from utils.lmdb_utils import load_from_lmdb +from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id +from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection + +logger = logging.getLogger(__name__) + + +def get_prepared_sequence_length(input_sequence: np.ndarray, model_configuration: dict) -> int: + """ + Returns the length of the input sequence after padding and resampling + """ + + initial_length = len(input_sequence) + take_every_nth = model_configuration["preprocessing"]["take_every_nth"] + padding_length = get_padding_length(model_configuration, resampled=False) + raw_length = initial_length + padding_length + adjusted = int(raw_length // take_every_nth) + return adjusted + + +def get_padding_length(model_configuration: dict, resampled: bool = True) -> int: + """ + Returns the length of the input sequence after padding and resampling + """ + take_every_nth = model_configuration["preprocessing"]["take_every_nth"] + padding_length = int(model_configuration["input_window_length"] * ( + 1 - model_configuration["preprocessing"]["min_input_length_fraction_for_padding"])) * ( + take_every_nth if not resampled else 1) + return padding_length + + +def get_number_of_windows(base_length: int, model_configuration: dict) -> int: + """ + Returns the number of windows for a given sequence length + """ + window_shift = model_configuration["preprocessing"]["window_shift"] + input_window_length = model_configuration["input_window_length"] + output_window_length = model_configuration["output_window_length"] + output_window_offset = model_configuration["output_window_offset"] + + if input_window_length > output_window_length + output_window_offset: + return (base_length - input_window_length) // window_shift + 1 + else: + return (base_length - output_window_length - output_window_offset) // window_shift + 1 + + +def get_collated_batch_for_key(sample_key: ObjectId | str, + model_configuration: dict, + start_cutoff: int = None, + end_cutoff: int = None, + lmdb_env=None) -> dict: + sample_batch_for_key = get_batch_for_key(sample_key, model_configuration, + start_cutoff=start_cutoff, + end_cutoff=end_cutoff, + lmdb_env=lmdb_env) + collated = model_configuration["collate_fn"](sample_batch_for_key) + return collated + + +def get_batch_for_key(key, + model_configuration: dict, + start_cutoff: int = None, + end_cutoff: int = None, + lmdb_env=None) -> np.ndarray: + """ + Get the batch for a given key from the lmdb database or compute it directly. + + The return has the same format as the batch produced by the batch_fn in the model configuration. + Args: + key: key to get the batch for, can also be object id from database + model_configuration: model configuration to use + start_cutoff: start cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points + end_cutoff: end cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points + lmdb_env: the lmdb environment to use, if None, the features are computed directly from the database + + Returns: + batch: batch for the given key, as returned by the batch_fn in the model configuration + + """ + if lmdb_env is None: + # compute and scale features, ignored features are handled internally + features = get_scaled_feature_for_key(key, model_configuration) + else: + # load features from lmdb + features = load_from_lmdb(lmdb_env, str(key)) + # ignored features are handled internally, thus we need to remove them here + ignored_features = model_configuration["feature_config"]["ignored_features"] if "ignored_features" in \ + model_configuration[ + "feature_config"] else [] + # add scaled versions of features + ignored_features = ignored_features + [feat + "_scaled" for feat in ignored_features] + for feature_set in features: + for feature_name in list(features[feature_set].keys()): + if feature_name in ignored_features: + del features[feature_set][feature_name] + + processed_chunk = process_chunk([key], [features], model_configuration) + + # use cutoff, if provided + if start_cutoff is not None or end_cutoff is not None: + + item_length = len(processed_chunk[0]["target_features"]) + + if start_cutoff is None: + start_cutoff = 0 + if end_cutoff is None: + end_cutoff = item_length + + # "normalize" cutoff to adjust for added padding and resampling + padding_length = get_padding_length(model_configuration, resampled=False) + take_every_nth = model_configuration["preprocessing"]["take_every_nth"] + # make sure that there is no padding added to the start cutoff, so that it includes the padding added to the sequence + start_cutoff_normalized = max(int(start_cutoff // take_every_nth), 0) + # end cutoff must be adjusted to include the padding added to the sequence + end_cutoff_normalized = min(int((end_cutoff + padding_length) // take_every_nth), item_length) + for feature_set in processed_chunk[0]: + if feature_set in model_configuration["feature_config"]["feature_sets"]: + processed_chunk[0][feature_set] = processed_chunk[0][feature_set][ + start_cutoff_normalized:end_cutoff_normalized + 1] + + if "batch_fn" in model_configuration and model_configuration["batch_fn"] is not None: + batch = model_configuration["batch_fn"](processed_chunk, model_configuration) + else: + batch = processed_chunk + if batch is None or len(batch) == 0: + raise ValueError(f"Batch is empty for key {key}") + return batch + + +def get_scaled_feature_for_key(key: str, + model_configuration: dict) -> tuple: + feature_config = model_configuration["feature_config"] + + user_cycles = list( + get_cycles_collection().find({"user_id": ObjectId(key)} | feature_config["filter_criteria"]).sort("starts_at", + 1)) + cycle_features = list() + for cycle in user_cycles: + features = get_features(cycle, feature_config) + cycle_features.append(features) + + # combine all cycles for the user + if len(cycle_features) == 0: + raise ValueError(f"No features found for key {key}") + features = combine_features(cycle_features, feature_config) + scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers") + scalers = load_scalers(scaler_dir) + # scale features + scaled_features = scale_item(features, scalers) + return scaled_features + + +def augment_items(identifiers: list, + items: list[dict], + feature_config: dict, + lmdb_env: lmdb.Environment = None) -> list: + """ + Augments the given sequences by attaching previous sequences + Args: + identifiers: identifiers of items for finding previous sequences + items: actual items + feature_config: feature config of the dataset + lmdb_env: lmdb environment for loading previous sequences, can be None, in this case the items are computed directly + + Returns: + list of augmented items + """ + + if "augmentation" not in feature_config or \ + feature_config["augmentation"]["use_augmentation"] is False: + return items + + augmented_items = list() + for i, item in enumerate(items): + item_id = identifiers[i] + try: + item_data = get_cycles_collection().find_one({"_id": ObjectId(item_id)}, {"user_id": 1, "starts_at": 1}) + item_user_id = item_data["user_id"] + item_starts_at = item_data["starts_at"] + except KeyError: + print(f"User ID not found for item {item_id}") + continue + previous_item_identifiers = list(get_cycles_collection().aggregate( + # make sure to filter by user first to significantly reduce the number of items + [ + { + "$match": { + "user_id": item_user_id, + "starts_at": {"$lt": item_starts_at}, + } + } + ] + feature_config["filter_criteria_pipeline"] + [ + { + "$sort": { + "starts_at": -1 + } + }, + { + "$project": { + "_id": 1, + "starts_at": 1 + } + } + ] + )) + max_lookback = feature_config["augmentation"]["max_lookback"] + previous_items = list() + for previous_item in previous_item_identifiers: + if len(previous_items) >= max_lookback: + break + previous_item_id = previous_item["_id"] + if lmdb_env is None: + # compute item + cycle = get_cycle_by_id(previous_item_id) + features = get_features(cycle, feature_config) + scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers") + scalers = load_scalers(scaler_dir) + previous_item = scale_item(features, scalers) + previous_items.append(previous_item) + else: + # take item from lmdb + try: + # make sure to parse object id + previous_item = load_from_lmdb(lmdb_env, str(previous_item_id)) + except KeyError: + continue + previous_items.append(previous_item) + + # merge items into one + feature_sets = feature_config["feature_sets"] + augmented_item = copy.deepcopy(item) + for feature_set in feature_sets: + if feature_set not in augmented_item: + continue + if "static" in feature_set: + continue + for feature_name in augmented_item[feature_set]: + if feature_name not in augmented_item[feature_set]: + continue + augmented_item[feature_set][feature_name] = np.concatenate( + [previous_item[feature_set][feature_name] for previous_item in previous_items] + + [augmented_item[feature_set][feature_name]]) + + augmented_items.append(augmented_item) + + return augmented_items + + +def process_chunk(ids: list, + chunk: list, + model_configuration: dict, + ignored_features: list = None, + pad_sequences: bool = True, + statics_as_list: bool = True): + """ + Creates a chunk of data as dataframe for training and inference of a tft model. + :param ids: list of ids of the data points in the chunk + :param chunk: list of data points + + :return: dataframe with the data points in the chunk + """ + + if ignored_features is None: + ignored_features = list() + + input_features = list(chunk[0]["target_features"].keys()) + feature_config = model_configuration["feature_config"] + + # variables for padding + padding_value = 0 + + # create dataframe from items in chunk + data = [] + for i, item in enumerate(chunk): + item_length = item["target_features"][input_features[0]].shape[0] + take_every_nth = model_configuration["preprocessing"]["take_every_nth"] + + data_item = dict() + for feature_set in feature_config["feature_sets"]: + if feature_set in item and len(item[feature_set]) > 0: + # fill in the data + source = item[feature_set] + features_to_use = [feature for feature in source.keys() if "_scaled" in feature] + feature_set_data = [source[feature] for feature in features_to_use] + + padding_length = get_padding_length(model_configuration, resampled=False) + + if len(feature_set_data) == 0: + raise ValueError(f"No features found for key {ids[i]} in feature set {feature_set}") + + if len(feature_set_data[0]) == 1: + # if the feature is constant, we need to repeat it for all time points + if statics_as_list: + if pad_sequences: + desired_length = item_length + padding_length + feature_set_data = np.array( + [np.full((desired_length,), val) for val in feature_set_data]).T[::take_every_nth] + else: + feature_set_data = np.array( + [np.full((item_length,), val) for val in feature_set_data]).T[::take_every_nth] + else: + feature_set_data = np.array(feature_set_data).flatten() + else: + if pad_sequences: + padding_values = np.full((padding_length, len(feature_set_data)), + [padding_value for vals in feature_set_data]).T + feature_set_data = np.concatenate([padding_values, np.array(feature_set_data)], axis=1) + # create bins of n values and apply mean + + # check, if special accumulation function has been specified + individual_feature_data = list() + for i, feature_to_use in enumerate(features_to_use): + try: + # get config for feature, make sure to replace scaled to get actual config + current_feature_config = \ + [x for x in feature_config[feature_set] if + x["name"] == feature_to_use.replace("_scaled", "")][0] + current_accumulation_fn = current_feature_config["accumulation_fn"] + except (KeyError, IndexError): + # default to mean + current_accumulation_fn = np.mean + + current_feature_data = apply_fn_to_bins(feature_set_data[i], take_every_nth, + current_accumulation_fn) + individual_feature_data.append(current_feature_data) + + # stack individual features + feature_set_data = np.stack(individual_feature_data).T + + if feature_set not in data_item: + data_item[feature_set] = feature_set_data + else: + data_item[feature_set] = np.concatenate([data_item[feature_set], feature_set_data], axis=1) + + data.append(data_item) + + return data + + +def apply_fn_to_bins(input_sequence: np.ndarray, + bin_size: int, + fn: callable) -> np.ndarray: + """ + Applies a function to the bins of the input sequence along the last dimension and returns the results. + Args: + input_sequence: sequence to apply the function to (1D or 2D) + bin_size: size of the bins + fn: function to apply to the bins + + Returns: + np.ndarray: result of the function applied to the bins + + """ + if input_sequence.ndim == 1: + input_sequence = input_sequence.reshape(1, -1) + dims = 1 + else: + dims = input_sequence.shape[-1] + + result = [] + for row in input_sequence: + row_result = [] + for start in range(0, len(row), bin_size): + end = min(start + bin_size, len(row)) + bin_slice = row[start:end] + row_result.append(fn(bin_slice)) + result.append(row_result) + + if dims > 1: + return np.array(result) + else: + return np.array(result[0]) + + +def produce_window_batches(data_chunk: list, + model_configuration: dict, + offsets: list | np.ndarray = None) -> list: + """ + Produces batches of windows from the data chunk. + Args: + data_chunk: data chunk + model_configuration: configuration to use + offsets: offsets used to determine the length of data to use, uses item[-offset:] of data + + Returns: + data: list of windows + """ + input_window_length = model_configuration["input_window_length"] + output_window_length = model_configuration["output_window_length"] + output_window_offset = model_configuration["output_window_offset"] + window_shift = model_configuration["preprocessing"]["window_shift"] + + if offsets is None: + offsets = np.zeros(len(data_chunk), dtype=int) + + data = list() + for i, item in enumerate(data_chunk): + item_length = item["target_features"][-offsets[i]:].shape[0] + total_window_length = input_window_length if input_window_length > output_window_length + output_window_offset \ + else output_window_length + output_window_offset + + if item_length < total_window_length: + continue + + # number of windows is defined by the input window length and output window length with offset + num_windows = get_number_of_windows(item_length, model_configuration) + for j in range(num_windows): + window = dict() + for key, value in item.items(): + if value is None: + window[key] = None + else: + if key == "target_features": + window[key] = value[-offsets[i]:][j * window_shift + output_window_offset: + j * window_shift + output_window_length + output_window_offset] + else: + window[key] = value[-offsets[i]:][j * window_shift:j * window_shift + input_window_length] + + data.append(window) + + return data + + +def produce_simple_batches(data_chunk: list, + model_configuration: dict, + offsets: list | np.ndarray = None) -> list: + """ + Produces batches without windowing. + Args: + data_chunk: data chunk + model_configuration: configuration to use + offsets: offsets used to determine the length of data to use, uses item[-offset:] of data + + Returns: + data: list of items + """ + pass + + +class LMDBIterableDataset(IterableDataset): + def __init__(self, + lmdb_env_path: str, + lmdb_keys: list[str], + model_configuration: dict, + batch_size: int = 32): + self.lmdb_path = lmdb_env_path + self.lmdb_env = None + self.lmdb_keys = lmdb_keys + self.key_subset = None + self.model_configuration = model_configuration + self.batch_size = batch_size + self.keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"]) + self.len = None + + try: + self.model_configuration["batch_fn"] + except KeyError: + raise KeyError("Batch function not found in model configuration") + + def set_key_subset(self, key_subset: list[str]): + """ + Set the key subset to use for the dataset. + Args: + key_subset: list of keys to use + """ + self.key_subset = key_subset + # reset length + self.len = None + + def get_length_of_data_subset(self, key_set: list[str]): + """ + Get the length of the data subset. + Args: + key_set: list of keys to use + """ + # run simplified version of __iter__ to get the length + + num_steps = 0 + i = 0 + num_batch_items = 0 + while True: + if num_batch_items >= self.batch_size: + num_steps += 1 + num_batch_items -= self.batch_size + else: + if i >= len(key_set): + if num_batch_items > 0: + num_steps += 1 + break + key = key_set[i] + i += 1 + + # try to fetch stats from key_stats + try: + current_key_stats = self.keys_stats["by_key"][key] + base_length = current_key_stats["item_length"] + take_every_nth = self.model_configuration["preprocessing"]["take_every_nth"] + padding_length = get_padding_length(self.model_configuration, resampled=False) + raw_length = base_length + padding_length + item_length = int(raw_length // take_every_nth) + except: + item = load_from_lmdb(self.lmdb_env, key) + item_length = get_prepared_sequence_length( + item["target_features"][list(item["target_features"])[0]], + self.model_configuration) + num_windows = get_number_of_windows(item_length, self.model_configuration) + num_batch_items += num_windows + + return num_steps + + def __iter__(self): + self.init_lmdb_env() + + # if no key subset is set, use all keys + if self.key_subset is None: + self.key_subset = self.lmdb_keys + + random.shuffle(self.key_subset) + + batch = list() + counter = 0 + collate_fn = self.model_configuration["collate_fn"] + while True: + if len(batch) >= self.batch_size: + if collate_fn is None: + yield batch[:self.batch_size] + else: + yield collate_fn(batch[:self.batch_size]) + batch = batch[self.batch_size:] + else: + if counter >= len(self.key_subset): + if len(batch) > 0: + yield collate_fn(batch) + break + key = self.lmdb_keys[counter] + counter += 1 + try: + current_batch = get_batch_for_key(key, self.model_configuration, lmdb_env=self.lmdb_env) + except ValueError: + # if betch is empty, try next + continue + batch.extend(current_batch) + + def init_lmdb_env(self): + if self.lmdb_env is None: + self.lmdb_env = lmdb.open(self.lmdb_path, + readonly=True, + lock=False, + readahead=False, + meminit=False) + + def __len__(self): + self.init_lmdb_env() + + # if no key subset is set, use all keys + if self.key_subset is None: + self.key_subset = self.lmdb_keys + + if self.len is None: + num_steps = self.get_length_of_data_subset(self.key_subset) + self.len = num_steps + + return self.len diff --git a/code/new_realtime/utils/dataset_creation.py b/code/new_realtime/utils/dataset_creation.py new file mode 100644 index 0000000..5a9ceff --- /dev/null +++ b/code/new_realtime/utils/dataset_creation.py @@ -0,0 +1,235 @@ +import os + +import lmdb +import numpy as np +import pandas as pd +import torch +from sklearn.preprocessing import StandardScaler, MinMaxScaler +import joblib + +from utils.lmdb_utils import load_from_lmdb, get_lmdb_keys + + +def get_features(cycle: dict, + feature_config: dict) -> dict: + """ + Computes the feature for a given cycle defined by the feature_config + :param cycle: cycle data as dictionary + :param feature_config: feature config as dictionary + :return: dict with feature according feature config + """ + features = {} + feature_sets = feature_config["feature_sets"] + for feature_set in feature_sets: + for feature_def in feature_config[feature_set]: + if feature_set not in features: + features[feature_set] = {} + + # skip features that are marked as ignored + if "ignored_features" in feature_config and feature_def["name"] in feature_config["ignored_features"]: + continue + + feature_return = feature_def["fn"](cycle=cycle) + if isinstance(feature_return, dict): + if len(feature_return) > 1: + for feature_name, feature in feature_return.items(): + features[feature_set][f"{feature_def['name']}_{feature_name}"] = feature + else: + features[feature_set][feature_def["name"]] = list(feature_return.values())[0] + else: + features[feature_set][feature_def["name"]] = feature_return + + # check, if all features have same length + if feature_set in features and len(features[feature_set]) > 0 and isinstance( + list(features[feature_set].values())[0], np.ndarray): + feature_lengths = [len(x) for x in features[feature_set].values()] + if len(set(feature_lengths)) > 1: + raise ValueError(f"Feature set {feature_set} has features of different lengths: {feature_lengths}") + + return features + + +def save_scalers(scalers: dict, + scaler_dir: str): + if not os.path.exists(scaler_dir): + os.makedirs(scaler_dir) + for feature_type in scalers: + for feature_name in scalers[feature_type]: + joblib.dump(scalers[feature_type][feature_name], + f"{scaler_dir}/{feature_name}.pkl") + + +def load_scalers(scaler_dir: str) -> dict: + scalers = dict() + for scaler_file in os.listdir(scaler_dir): + if scaler_file.endswith(".pkl"): + feature_name = scaler_file.replace(".pkl", "") + scalers[feature_name] = joblib.load(f"{scaler_dir}/{scaler_file}") + return scalers + + +def get_scalers_for_model(model_configuration: dict): + """ + Get the scalers for the features of a model configuration + Args: + model_configuration: configuration of the model + + Returns: + scalers: scalers for the model + + """ + + scaler_dir = os.path.join(model_configuration["feature_config"]["dataset_dir"], "scalers") + scalers = load_scalers(scaler_dir) + return scalers + + +def get_feature_values(feature_type: str | None, + feature_name: str, + env: lmdb.Environment, + keys: list = None) -> list: + """ + Get the values of a feature from the lmdb dataset for all keys + :param feature_type: type of feature, can be None, then the first feature with the given name will be used + :param feature_name: feature name to extract + :param env: lmdb environment + :param keys: keys to extract the feature from, if None, all keys will be used / fetched from the database + :return: list with feature values + """ + if keys is None: + keys = get_lmdb_keys(env) + feature_values = [] + for key in keys: + data = load_from_lmdb(env, key) + if feature_type is None: + for feature_type, feature_data in data.items(): + if feature_name in feature_data: + feature_values.append(feature_data[feature_name]) + break + else: + for current_feature_name, feature_data in data[feature_type].items(): + # catch sub features that have been prefixed with the feature name + if current_feature_name.startswith(feature_name): + feature_values.append(feature_data) + break + return feature_values + + +def train_scalers(feature_type: str, + feature_name: str, + scaler_type, + sample, + env) -> dict: + scalers = dict() + all_features = sample[feature_type] + individual_feature_names = list() + for individual_feature_name in all_features: + if individual_feature_name.startswith(feature_name) and not individual_feature_name.endswith("_scaled"): + individual_feature_names.append(individual_feature_name) + + for individual_feature_name in individual_feature_names: + feature_values = get_feature_values(feature_type, individual_feature_name, env) + if len(feature_values) == 0: + raise ValueError(f"No feature values found for feature {individual_feature_name}") + if isinstance(feature_values[0], list) or isinstance(feature_values[0], np.ndarray): + features_reshaped = np.concatenate(feature_values).reshape(-1, 1) + else: + features_reshaped = np.array(feature_values).reshape(-1, 1) + del feature_values + + if scaler_type is not None: + scaler = scaler_type() + scaler.fit(features_reshaped) + else: + scaler = None + + scalers[individual_feature_name] = scaler + return scalers + + +def scale_item(data: dict, + scalers: dict): + data_format = dict() + for feature_set in data: + if feature_set not in data_format: + data_format[feature_set] = dict() + for feature_name in data[feature_set]: + if feature_name.endswith("_scaled"): + continue + data_format[feature_set][feature_name] = data[feature_set][feature_name] + + for feature_set in data_format: + for feature_name in data_format[feature_set]: + feature_values = data[feature_set][feature_name] + if scalers[feature_name] is not None: + if isinstance(feature_values, list) or isinstance(feature_values, np.ndarray): + scaled_feature_values = scalers[feature_name].transform( + feature_values.reshape(-1, 1)).flatten() + else: + scaled_feature_values = scalers[feature_name].transform( + np.array(feature_values).reshape(-1, 1)).flatten() + else: + scaled_feature_values = feature_values + data[feature_set][f"{feature_name}_scaled"] = scaled_feature_values + + return data + + +def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor, + feature_names: np.ndarray | list | str, + scalers: dict) -> np.ndarray: + """ + Inverse scales the input features, can handle single and multi feature input + Args: + input_feature: input feature as a numpy array + feature_names: names of input features as reference for scalers + scalers: dict of feature scalers + + Returns: + scaled input features as numpy array + """ + + if isinstance(feature_names, str): + feature_names = [feature_names] + + if isinstance(input_feature, np.ndarray) or isinstance(input_feature, torch.Tensor): + input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1 + elif isinstance(input_feature, float) or isinstance(input_feature, int): + # handle case, where input is scalar + input_scaled = scalers[feature_names[0]].inverse_transform([[input_feature]])[0][0] + return input_scaled + elif isinstance(input_feature, list): + input_feature = np.array(input_feature) + input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1 + else: + raise ValueError(f"Unsupported input type: {type(input_feature)}") + + # reshape, if input is one dimensional + if input_dim == 1: + input_feature = input_feature.reshape(-1, 1) + + input_scaled = input_feature.copy() + for i in range(input_dim): + input_scaled[:, i] = scalers[feature_names[i]].inverse_transform(np.array([input_scaled[:, i]])) + + return input_scaled + + +def combine_features(cycles: list[dict], feature_config: dict) -> dict: + combined_features = dict() + for feature_set in feature_config["feature_sets"]: + combined_features[feature_set] = dict() + if feature_set not in cycles[0]: + continue + features_in_set = cycles[0][feature_set].keys() + for feature in features_in_set: + feature_values = [cycles[i][feature_set][feature] for i in range(len(cycles))] + if isinstance(feature_values[0], dict): + for key in feature_values[0].keys(): + feature_array = np.concatenate([feature_values[i][key] for i in range(len(feature_values))]) + combined_features[feature_set][key] = feature_array + else: + feature_array = np.concatenate(feature_values) + combined_features[feature_set][feature] = feature_array + + return combined_features diff --git a/code/new_realtime/utils/dataset_utils.py b/code/new_realtime/utils/dataset_utils.py new file mode 100644 index 0000000..fc45b3d --- /dev/null +++ b/code/new_realtime/utils/dataset_utils.py @@ -0,0 +1,36 @@ +import os +import pickle + + +def get_dataset_path(lmdb_base_dir: str, + dataset_name: str) -> str: + """ + Get the path to the dataset in the lmdb directory + Args: + lmdb_base_dir: base directory of the lmdb dataset + dataset_name: name of the dataset + + Returns: + path to the dataset + """ + + dataset_path = os.path.join(lmdb_base_dir, dataset_name) + if not os.path.exists(dataset_path): + raise FileNotFoundError(f"Dataset {dataset_name} not found in {lmdb_base_dir}") + return dataset_path + + +def load_key_stats(lmdb_dir: str) -> dict: + """ + Load key statistics from the LMDB database. + :param lmdb_dir: Directory of the LMDB database. + :return: Dictionary with statistics. + """ + key_stats_path = f"{lmdb_dir}/key_stats.pickle" + + if not os.path.exists(key_stats_path): + return dict() + + with open(key_stats_path, "rb") as f: + key_stats = pickle.load(f) + return key_stats diff --git a/code/new_realtime/utils/evaluation.py b/code/new_realtime/utils/evaluation.py new file mode 100644 index 0000000..54f68bc --- /dev/null +++ b/code/new_realtime/utils/evaluation.py @@ -0,0 +1,266 @@ +import lmdb +import numpy as np +import sklearn +import torch +from torch import nn +from tqdm import tqdm + +from utils.data_utils import get_collated_batch_for_key, get_padding_length +from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature +from utils.dataset_utils import load_key_stats + +from vsm_datascience_common import constants + + +def evaluate_model(model_configuration: dict, + training_configuration: dict, + test_ids: list, + evaluation_functions: list) -> dict: + dataset_dir = model_configuration["feature_config"]["dataset_dir"] + lmdb_env = lmdb.open(dataset_dir, readonly=True) + # get computation rank + if torch.distributed.is_initialized(): + local_rank = torch.distributed.get_rank() + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + else: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + load_fn = model_configuration["model_load_fn"] + model = load_fn(model_configuration=model_configuration, + training_configuration=training_configuration, + sample_key=test_ids[0], + device=device, + lmdb_env=lmdb_env) + # load key stats to retrieve individual cycles + keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"]) + + batch_size = training_configuration["batch_size"] + predict_fn = model_configuration["predict_fn"] + actual_fn = model_configuration["actual_fn"] + + target_features = [x["name"] for x in model_configuration["feature_config"]["target_features"]] + ignored_features = model_configuration["feature_config"]["ignored_features"] + used_targets = [x for x in target_features if x not in ignored_features] + scalers = get_scalers_for_model(model_configuration) + + errors = dict() + for test_id in tqdm(test_ids): + # fetch stats for key + current_key_stats = keys_stats["by_key"][test_id] + cycle_stats = current_key_stats["cycle_stats"] + for i in range(len(cycle_stats)): + # compute cutoffs to isolate current cycle + current_cycle_start_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i]]) + current_cycle_end_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i + 1]]) + try: + batch = get_collated_batch_for_key(test_id, model_configuration, + start_cutoff=current_cycle_start_cutoff, + end_cutoff=current_cycle_end_cutoff, + lmdb_env=lmdb_env) + except ValueError: + print(f"Skipping {test_id}") + continue + if batch is None: + print(f"Empty batch, skipping {test_id}") + continue + + preds = predict_fn(model, batch, + batch_size=batch_size, + device=device) + actuals = actual_fn(batch) + num_outputs = preds.shape[-1] if len(preds.shape) > 1 else 1 + scaled_preds = list() + for j in range(num_outputs): + if isinstance(preds, torch.Tensor): + preds = preds.cpu().numpy() + + if len(preds.shape) == 3: + output = preds[:, 0, j].squeeze() + else: + output = preds[:, j] if num_outputs > 1 else preds + + if "loss_functions" in training_configuration and len(training_configuration["loss_functions"]) > j: + if isinstance(training_configuration["loss_functions"][j], nn.BCEWithLogitsLoss): + output = torch.sigmoid(torch.tensor(output)).numpy() + + scaled_output = inverse_scale_feature(output, + used_targets[j], + scalers) + + scaled_preds.append(scaled_output) + + scaled_actuals = list() + for j in range(num_outputs): + if len(actuals.shape) == 3: + output = actuals[:, 0, j].squeeze() + else: + output = actuals[:, j] if num_outputs > 1 else actuals + if isinstance(output, torch.Tensor): + output = output.cpu().numpy() + scaled_output = inverse_scale_feature(output, + used_targets[j], + scalers).ravel() + + scaled_actuals.append(scaled_output) + + for eval_fn in evaluation_functions: + if eval_fn is not None: + eval_fn_name = eval_fn["name"] + eval_function = eval_fn["eval_fn"] + eval_fn_index = eval_fn["input_index"] + # skip error fn if actuals are nan, since they are ignored + if any(np.isnan(scaled_actuals[eval_fn_index])): + continue + error = eval_function(scaled_preds[eval_fn_index], scaled_actuals[eval_fn_index]) + + if np.isnan(error): + # skip if error is nan + continue + + if eval_fn_name not in errors: + errors[eval_fn_name] = dict() + + if f"after_{i}" not in errors[eval_fn_name]: + errors[eval_fn_name][f"after_{i}"] = list() + + errors[eval_fn_name][f"after_{i}"].append(error) + + for eval_fn in evaluation_functions: + if eval_fn is not None: + eval_fn_name = eval_fn["name"] + accumulation_fn = eval_fn["accumulation_fn"] + for key in errors[eval_fn_name]: + if len(errors[eval_fn_name][key]) == 0: + errors[eval_fn_name][key] = np.nan + else: + errors[eval_fn_name][key] = accumulation_fn(errors[eval_fn_name][key]) + return errors + + +def pre_ov_error(preds, actuals, *args, **kwargs): + if any(np.isnan(actuals)): + return np.nan + # get ov day index + try: + ov_day_index = np.where(actuals == 0)[0][0] + except IndexError: + # no ov day in actuals + return np.nan + + # get pre ov predictions + pre_ov_preds = preds[:ov_day_index] + pre_ov_actuals = actuals[:ov_day_index] + + if len(pre_ov_preds) == 0 or len(pre_ov_actuals) == 0: + return np.nan + + # calculate error + error = sklearn.metrics.mean_absolute_error(pre_ov_preds, pre_ov_actuals) + return error + + +def after_ov_error(preds, actuals, *args, **kwargs): + if any(np.isnan(actuals)): + return np.nan + # get ov day index + try: + ov_day_index = np.where(actuals == 0)[0][0] + except IndexError: + # no ov day in actuals + return np.nan + # get pre ov predictions + after_ov_preds = preds[ov_day_index:] + after_ov_actuals = actuals[ov_day_index:] + + if len(after_ov_preds) == 0 or len(after_ov_actuals) == 0: + return np.nan + + # calculate error + error = sklearn.metrics.mean_absolute_error(after_ov_preds, after_ov_actuals) + return error + + +def ov_error(preds, actuals, model_configuration, *args, **kwargs): + if any(np.isnan(actuals)): + return np.nan + # get ov day index + try: + ov_day_index = np.where(actuals == 0)[0][0] + except IndexError: + # no ov day in actuals + return np.nan + + # get predicted ov index + try: + pred_ov_index = np.where(preds >= 0)[0][0] + except IndexError: + # no ov day in actuals + return np.nan + + # calculate error + error = abs(pred_ov_index - ov_day_index) + + # scale error to account for step size + step_size = model_configuration["preprocessing"]["window_shift"] + measurements_per_day = constants.MEASUREMENTS_PER_DAY + downsampling_factor = model_configuration["preprocessing"]["take_every_nth"] + shift_hour_factor = 24 // (measurements_per_day // downsampling_factor) * step_size + error_in_days = error * shift_hour_factor // 24 + + return error_in_days + + +def day_relative_to_ov_error(preds: np.ndarray | torch.Tensor, + actual: np.ndarray | torch.Tensor, + day_relative_to_ov: int, + model_configuration: dict) -> float: + """ + Calculate the error of the model predictions relative to the ov day + Args: + preds: predictions + actual: actual values + day_relative_to_ov: day relative to ov day + model_configuration: model configuration + + Returns: + error: error of the model predictions relative to the ov day + """ + + if isinstance(preds, torch.Tensor): + preds = preds.cpu().numpy() + + if isinstance(actual, torch.Tensor): + actual = actual.cpu().numpy() + + if any(np.isnan(actual)): + return np.nan + + # get ov day index + try: + ov_day_index = np.where(actual == 0)[0][0] + except IndexError: + # no ov day in actuals + return np.nan + + # get index offset factor -> how much time between each step + step_size = model_configuration["preprocessing"]["window_shift"] + measurements_per_day = constants.MEASUREMENTS_PER_DAY + downsampling_factor = model_configuration["preprocessing"]["take_every_nth"] + shift_hour_factor = 24 // (measurements_per_day // downsampling_factor) * step_size + + index_offset = int(day_relative_to_ov * (shift_hour_factor // 24)) + + if ov_day_index + index_offset >= preds.shape[0]: + relative_day_pred = preds[-1] if len(preds.shape) > 1 else preds[-1] + elif ov_day_index + index_offset < 0: + relative_day_pred = preds[0][0] if len(preds.shape) > 1 else preds[0] + else: + relative_day_pred = preds[ov_day_index + index_offset][0] if len(preds.shape) > 1 else \ + preds[ov_day_index + index_offset][0] + + relative_day_actual = day_relative_to_ov + + # calculate error + error = np.abs(relative_day_pred - relative_day_actual) + return error diff --git a/code/new_realtime/utils/feature_functions.py b/code/new_realtime/utils/feature_functions.py new file mode 100644 index 0000000..1a25bf8 --- /dev/null +++ b/code/new_realtime/utils/feature_functions.py @@ -0,0 +1,372 @@ +from datetime import timedelta +from typing import Callable + +import numpy as np +from vsm_datascience_common import constants +from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection +from vsm_datascience_common.cycles.sequences import get_timestamps, get_values + +from utils.smoothing import get_curve_composition + + +def get_day_of_week_encoded(cycle: dict, shift: int = 0) -> dict: + timestamps = get_timestamps(cycle) + weekdays = np.array([x.weekday() for x in timestamps]) + sine_encoded = np.sin(weekdays * (2 * np.pi / 7)) + cosine_encoded = np.cos(weekdays * (2 * np.pi / 7)) + return { + "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None], + "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None] + } + + +def get_hour_of_day_encoded(cycle: dict, shift: int = 0) -> dict: + timestamps = get_timestamps(cycle) + hours = np.array([x.hour for x in timestamps]) + sine_encoded = np.sin(hours * (2 * np.pi / 24)) + cosine_encoded = np.cos(hours * (2 * np.pi / 24)) + return { + "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None], + "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None] + } + + +def get_month_of_year_encoded(cycle: dict, shift: int = 0) -> dict: + timestamps = get_timestamps(cycle) + months = np.array([x.month for x in timestamps]) + sine_encoded = np.sin(months * (2 * np.pi / 12)) + cosine_encoded = np.cos(months * (2 * np.pi / 12)) + return { + "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None], + "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None] + } + + +def get_hours_from_start(cycle: dict, shift: int = 0) -> np.ndarray: + timestamps = get_timestamps(cycle) + return np.arange(len(timestamps))[shift if shift > 0 else 0:shift if shift < 0 else None] + + +def get_temperature(cycle: dict, shift: int = 0) -> np.ndarray: + return get_values(cycle)[shift if shift > 0 else 0:shift if shift < 0 else None] + + +def get_fertility_curve(cycle_length: int, + ov_index: int, + fertility_base_curve: np.ndarray | list, + base_curve_offset: int) -> np.ndarray: + fert_prob = np.zeros((cycle_length,)) + if ov_index + base_curve_offset < 0: + length_in_cycle = len(fertility_base_curve) + ov_index + base_curve_offset + fert_prob[:length_in_cycle] = fertility_base_curve[-length_in_cycle:] + elif ov_index + base_curve_offset + len(fertility_base_curve) > cycle_length: + length_in_cycle = cycle_length - ov_index - base_curve_offset + fert_prob[ov_index + base_curve_offset:] = fertility_base_curve[:length_in_cycle] + else: + fert_prob[ov_index + base_curve_offset: + ov_index + base_curve_offset + len(fertility_base_curve)] = fertility_base_curve + return fert_prob + + +def get_fertility_probability_base_curve() -> (np.ndarray, int): + """ + This function returns the base curve for the fertility probability as well as the offset to the ovulation day in measurements + :return: numpy array with the base curve resampled to the number of measurements per day + """ + fertilization_chance_curve = [0, 0.07, 0.12, 0.25, 0.3, 0.18, 0] + offset_to_ov = -6 + + new_indices = np.arange(len(fertilization_chance_curve) * constants.MEASUREMENTS_PER_DAY) + fertility_chance_curve_resampled = np.interp(new_indices, + np.linspace(0, len(new_indices), num=len(fertilization_chance_curve)), + fertilization_chance_curve) + ov_offset_resampled = offset_to_ov * constants.MEASUREMENTS_PER_DAY + return fertility_chance_curve_resampled, ov_offset_resampled + + +def get_fertility_probability(cycle: dict, shift: int = 0) -> np.ndarray: + # check, if biphasic + if "classification_results" in cycle and "results" in cycle["classification_results"][0]: + if cycle["classification_results"][0]["results"]["predicted_class"] != "biphasic": + return np.zeros((len(get_values(cycle)) - abs(shift, ))) + else: + return np.zeros((len(get_values(cycle)) - abs(shift, ))) + + if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]: + try: + ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"] + timestamps = get_timestamps(cycle) + ov_index = np.where(timestamps >= ov_timestamp)[0][0] + fertility_chance_curve_resampled, ov_offset_resampled = get_fertility_probability_base_curve() + return get_fertility_curve(len(get_values(cycle)), ov_index, fertility_chance_curve_resampled, + ov_offset_resampled)[shift if shift > 0 else 0:shift if shift < 0 else None] + except: + return np.zeros((len(get_values(cycle)) - abs(shift, ))) + else: + return np.zeros((len(get_values(cycle)) - abs(shift, ))) + + +def get_ov_over_probability(cycle: dict, shift: int = 0) -> np.ndarray: + """ + Get the ovulation over probability for a cycle + :param cycle: cycle data + :param shift: shift to apply + :return: ovulation over probability + """ + if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]: + try: + ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"] + timestamps = get_timestamps(cycle) + ov_index = np.where(timestamps > ov_timestamp)[0][0] + ov_over_probability = np.zeros(len(get_values(cycle))) + ov_over_probability[ov_index:] = 1 + return ov_over_probability[shift if shift > 0 else 0:shift if shift < 0 else None] + except: + return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None] + else: + return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None] + + +def get_days_relative_to_ov(cycle: dict) -> np.ndarray: + """ + Get the time in days until the next OV event. + """ + timestamps = get_timestamps(cycle) + if ("ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]) and ( + "classification_results" in cycle and "results" in cycle["classification_results"][0] and + cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"): + try: + ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"] + ov_index = np.where(timestamps > ov_timestamp)[0][0] + ov_timestamp = timestamps[ov_index] + days_until_ov = np.empty(len(timestamps)) + for i in range(len(timestamps)): + current_day = timestamps[i] + day_diff = ov_timestamp - current_day + if day_diff.days < -1: + pass + # days_until_ov[i] = max(day_diff.days -1) + # use - days, as days relative to event + days_until_ov[i] = - day_diff.days + + return np.array(days_until_ov) + except: + pass + + return np.full((len(timestamps),), np.nan) + + +def get_ov_day(cycle: dict) -> np.ndarray: + """ + Get the ovulation day for a cycle + :param cycle: cycle data + :return: ovulation day + """ + values = get_values(cycle) + seq_len = len(values) + is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \ + cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic" + has_ov_results = "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0] + if is_biphasic and has_ov_results: + try: + ov_day = cycle["ov_detection_results"][0]["results"]["ovulation_day"] + return np.full((seq_len,), ov_day) + except: + pass + + return np.full((seq_len,), np.nan) + + +def get_is_biphasic(cycle: dict) -> np.ndarray: + """ + Get the is_biphasic flag for a cycle + :param cycle: cycle data + :return: is_biphasic flag + """ + values = get_values(cycle) + seq_len = len(values) + is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \ + cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic" + if is_biphasic: + return np.full((seq_len,), 1) + + return np.full((seq_len,), 0) + + +def get_curve_composition_as_features(cycle: dict, shift: int = 0) -> dict: + composition = get_curve_composition(get_values(cycle)) + return { + "trend": composition[0][shift if shift > 0 else 0:shift if shift < 0 else None], + "seasonal": composition[1][shift if shift > 0 else 0:shift if shift < 0 else None], + "residual": composition[2][shift if shift > 0 else 0:shift if shift < 0 else None], + "smoothed": composition[3][shift if shift > 0 else 0:shift if shift < 0 else None] + } + + +def get_rolling_average_with_padding(cycle: dict, + shift: int = 0, + rolling_average_window_length: int = constants.MEASUREMENTS_PER_DAY) -> np.ndarray: + """ + Get the rolling average of the temperature values with padding, so that the rolling average has the same length as the + original temperature values + Args: + cycle (dict): cycle data + shift (int): shift to apply + rolling_average_window_length (int): window length for the rolling average, default is one day (288 measurements) + + Returns: + np.ndarray: rolling average of the temperature values with padding + + """ + + values = get_values(cycle) + values_padded = np.concatenate( + [np.full((constants.MEASUREMENTS_PER_DAY,), np.mean(values[:constants.MEASUREMENTS_PER_DAY])), values]) + rolling_average = np.convolve(values_padded, np.ones(rolling_average_window_length) / rolling_average_window_length, + mode="valid")[-len(values):] + return rolling_average[shift if shift > 0 else 0:shift if shift < 0 else None] + + +def get_window_fn(cycle: dict, + window_size: int, + fn: Callable) -> np.ndarray: + """ + Apply a window function to a time series. + :param cycle: cycle data + :param window_size: window size + :return: array with minimum value in window for each value in input array + """ + values = get_values(cycle) + values_padded = np.concatenate([np.full(window_size - 1, values[0]), values]) + + windows = np.lib.stride_tricks.sliding_window_view(values_padded, window_shape=window_size) + return fn(windows) + + +def get_cycle_length_stats(cycle: dict) -> dict: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycles_dates = list(get_cycles_collection().find({ + "user_id": user_id, "ends_at": {"$exists": True}, + "starts_at": {"$lt": cycle["starts_at"]}, + }, + {"starts_at": 1, "ends_at": 1})) + if len(cycles_dates) == 0: + # use the average cycle length as fallback + average_cycle_length = 37.13 + cycle_length_std = 12.42 + else: + cycle_lengths = [x["ends_at"] - x["starts_at"] for x in cycles_dates] + average_cycle_length = np.mean(cycle_lengths).total_seconds() / (60 * 60 * 24) + cycle_length_std = np.std([x.total_seconds() for x in cycle_lengths]) / (60 * 60 * 24) + + return { + "average_cycle_length": np.full((cycle_length,), average_cycle_length), + "cycle_length_std": np.full((cycle_length,), cycle_length_std), + } + + +def get_average_ovulation_day(cycle: dict) -> float: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycles_dates = list( + get_cycles_collection().find({ + "user_id": user_id, + "ov_detection_results.0.results": {"$exists": True}, + "starts_at": {"$lt": cycle["starts_at"]}, + }, + {"ov_detection_results": 1})) + if len(cycles_dates) == 0: + # use the average ovulation day as fallback + average_ovulation_day = 18.9 + else: + ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates] + average_ovulation_day = np.mean(ovulation_days) + return np.full((cycle_length,), average_ovulation_day) + + +def get_ovulation_std(cycle: dict) -> float: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycles_dates = list( + get_cycles_collection().find({ + "user_id": user_id, + "ov_detection_results.0.results": {"$exists": True}, + "starts_at": {"$lt": cycle["starts_at"]}, + }, + {"ov_detection_results": 1})) + if len(cycles_dates) == 0: + # use the average ovulation std as fallback + ovulation_std = 4.07 + else: + ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates] + ovulation_std = np.std(ovulation_days) + return np.full((cycle_length,), ovulation_std) + + +def get_biphasic_fraction(cycle: dict) -> float: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycle_dates = list( + get_cycles_collection().find({ + "user_id": user_id, + "classification_results.0.results": {"$exists": True}, + "starts_at": {"$lt": cycle["starts_at"]}, + }, + {"classification_results": 1})) + if len(cycle_dates) == 0: + # use the average cycle length as fallback + biphasic_fraction = 0.9449 + else: + biphasic = [x for x in cycle_dates if + x["classification_results"][0]["results"]["predicted_class"] == "biphasic"] + biphasic_fraction = len(biphasic) / len(cycle_dates) + return np.full((cycle_length,), biphasic_fraction) + + +def get_num_cycles(cycle: dict) -> int: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycles_dates = list(get_cycles_collection().find({"user_id": user_id, + "ends_at": {"$exists": True}, + "starts_at": {"$lt": cycle["starts_at"]} + }, + {"starts_at": 1, "ends_at": 1})) + num_cycles = len(cycles_dates) + return np.full((cycle_length,), num_cycles) + + +def get_average_temperatures(cycle: dict) -> dict: + user_id = cycle["user_id"] + cycle_length = len(get_values(cycle)) + cycles_dates = list( + get_cycles_collection().find({"user_id": user_id, "ov_detection_results.0.results": {"$exists": True}}, + {"starts_at": 1, "ov_detection_results": 1, "measurements": 1})) + + if len(cycles_dates) == 0: + return { + "pre_ov_temperatures": np.full((cycle_length,), 37.07), + "post_ov_temperatures": np.full((cycle_length,), 37.37), + } + + pre_ov_temperatures = [] + post_ov_temperatures = [] + for x in cycles_dates: + ov_detection = x["ov_detection_results"][0]["results"] + values = get_values(x) + timestamps = get_timestamps(x) + cutoff_day = x["starts_at"] + timedelta(days=ov_detection["ovulation_day"]) + cutoff_index = np.where(np.array(timestamps) >= cutoff_day)[0][0] + pre_ov_temperatures.append(values[:cutoff_index]) + post_ov_temperatures.append(values[cutoff_index:]) + + pre_ov_temperatures = np.concatenate(pre_ov_temperatures) + post_ov_temperatures = np.concatenate(post_ov_temperatures) + + pre_ov_temperatures = np.mean(pre_ov_temperatures) + post_ov_temperatures = np.mean(post_ov_temperatures) + + return { + "pre_ov_temperatures": np.full((cycle_length,), pre_ov_temperatures), + "post_ov_temperatures": np.full((cycle_length,), post_ov_temperatures), + } diff --git a/code/new_realtime/utils/lmdb_utils.py b/code/new_realtime/utils/lmdb_utils.py new file mode 100644 index 0000000..bdb8ef4 --- /dev/null +++ b/code/new_realtime/utils/lmdb_utils.py @@ -0,0 +1,119 @@ +import pickle + +import pandas as pd + + +def save_to_lmdb(env, key, dataset): + """ + Saves the given dataset to an LMDB environment with the given key. + :param env: LMDB environment + :param key: key to save the dataset to + :param dataset: tuple of pandas dataframes + """ + + with env.begin(write=True) as txn: + txn.put(key.encode('ascii'), pickle.dumps(dataset)) + + +def load_from_lmdb(env, key): + """ + Loads a dataset from an LMDB environment with the given key. + :param env: LMDB environment + :param key: key of the dataset to load + :return: key and tuple of pandas dataframes (input, context, output) + """ + with env.begin(write=False) as txn: + try: + data = pickle.loads(txn.get(key.encode('ascii'))) + return data + except TypeError: + raise KeyError(key) + + +def delete_from_lmdb(env, key): + """ + Deletes a dataset from an LMDB environment with the given key. + :param env: LMDB environment + :param key: key of the dataset to delete + """ + with env.begin(write=True) as txn: + txn.delete(key.encode('ascii')) + + +def clear_lmdb(env): + """ + Clears all datasets from an LMDB environment. + :param env: LMDB environment + """ + with env.begin(write=True) as txn: + cursor = txn.cursor() + for key, value in cursor: + txn.delete(key) + + +def lmdb_dataset_generator(env): + """ + Generator function to yield datasets from an LMDB environment. + :param env: LMDB environment + :return: generator + """ + with env.begin(write=False) as txn: + cursor = txn.cursor() + for key, value in cursor: + data = pickle.loads(value) + yield data + + +def lmdb_contains(env, substring) -> bool: + """ + Checks if the given substring is contained in any of the keys of the LMDB environment. + :param env: LMDB environment + :param substring: substring to search for + :return: boolean + """ + with env.begin(write=False) as txn: + cursor = txn.cursor() + for key, value in cursor: + if substring in key.decode('ascii'): + return True + return False + + +def lmdb_substring_key_search(env, substring): + """ + Searches for keys in the LMDB environment that contain the given substring. + :param env: LMDB environment + :param substring: substring to search for + :return: list of keys + """ + keys = [] + with env.begin(write=False) as txn: + cursor = txn.cursor() + for key, value in cursor: + if substring in key.decode('ascii'): + keys.append(key) + return keys + + +def get_lmdb_keys(env, limit: int = None): + """ + Get all keys in the LMDB environment. + :param env: LMDB environment + :param limit: maximum number of keys to return + :return: list of keys + """ + with env.begin(write=False) as txn: + with txn.cursor() as cursor: + keys = [key.decode("ascii") for key in cursor.iternext(keys=True, values=False)] + + return keys + + +def get_lmdb_keyspace_size(env): + """ + Get the number of keys in the LMDB environment. + :param env: LMDB environment + :return: number of keys + """ + with env.begin(write=False) as txn: + return txn.stat()['entries'] diff --git a/code/new_realtime/utils/loss_functions.py b/code/new_realtime/utils/loss_functions.py new file mode 100644 index 0000000..79f116e --- /dev/null +++ b/code/new_realtime/utils/loss_functions.py @@ -0,0 +1,118 @@ +import torch +from torch import nn +import torch.nn.functional as F +import numpy as np + + +class QuantileLoss(nn.Module): + def __init__(self, config): + super().__init__() + self.register_buffer('q', torch.tensor(config.quantiles)) + + def forward(self, predictions, targets): + diff = predictions - targets.unsqueeze(-1).expand(-1, -1, -1, predictions.shape[-1]) + ql = (1 - self.q) * F.relu(diff) + self.q * F.relu(-diff) + losses = ql.view(-1, ql.shape[-1]).mean(0) + return losses + + +def qrisk(pred, tgt, quantiles): + diff = pred - tgt + ql = (1 - quantiles) * np.clip(diff, 0, float('inf')) + quantiles * np.clip(-diff, 0, float('inf')) + losses = ql.reshape(-1, ql.shape[-1]) + normalizer = np.abs(tgt).mean() + risk = 2 * losses / normalizer + return risk.mean(0) + + +def weighted_bce_loss_fn(): + bce = nn.BCEWithLogitsLoss(reduction="none") # Use reduction="none" to get per-element loss + use_class_weights = False + time_weight_factor = 1 + + def _loss_fn(y_pred, y_true): + loss_list = [] # To accumulate per-sample loss values + for i in range(len(y_pred)): + current_pred = y_pred[i] + current_true = y_true[i] + loss_raw = bce(current_pred, current_true) + + if use_class_weights: + # Calculate class weights using torch operations + try: + factor = torch.max(current_true) / torch.mean(current_true) / 2 + non_zero_indices = current_true != 0 + class_weights = torch.ones_like(current_pred) + class_weights[non_zero_indices] = factor + except Exception: + class_weights = torch.ones_like(current_pred) + else: + class_weights = torch.ones_like(current_pred) + + # Calculate time weights using torch operations + try: + # Find indices where the event occurs + event_indices = torch.where(current_true == 1)[0] + if len(event_indices) > 0: + first_event_index = event_indices[0].item() + time_weights = torch.ones_like(current_pred) + # Weight the loss before the first event higher if needed + time_weights[:first_event_index] = time_weight_factor + else: + time_weights = torch.ones_like(current_pred) + except Exception: + time_weights = torch.ones_like(current_pred) + + # Apply the weights to the raw loss + loss = loss_raw * class_weights * time_weights + # Optionally, take the mean over the time dimension + loss_list.append(loss.mean()) + + # Aggregate the loss for the entire batch + return torch.stack(loss_list).mean() + + return _loss_fn + + +def weighted_mse_prob_loss_fn(): + mse = nn.MSELoss(reduce=False) + after_event_factor = 1 + use_weighted = False + + def _loss_fn(y_pred, y_true): + nonlocal mse, after_event_factor + loss_list = list() + for i in range(len(y_pred)): + current_pred = y_pred[i] + current_true = y_true[i] + loss_raw = mse(current_pred, current_true) + # calculate factor to use for weighting, only works with probability values + if use_weighted: + try: + factor = torch.max(current_true) / torch.mean(current_true) / 2 + non_zero_indices = current_true != 0 + class_weights = np.ones(len(current_true)) + class_weights[non_zero_indices] = factor + except: + # if there are no non-zero indices, set all weights to 1 + class_weights = np.ones(len(current_true)) + else: + # if we don't want to weight the loss, set all weights to 1 + class_weights = np.ones(len(current_true)) + + # calculate time weight, here, after the event, the loss should be higher + try: + last_event_index = np.where(current_true != 0)[-1] + time_weights = np.zeros(len(current_true)) + # here, we want to weight the loss after the 'fertility curve' higher, so that we don't get the signal too late + time_weights[last_event_index:] = after_event_factor + except: + # if there are no non-zero indices, set all weights to 1 + time_weights = np.ones(len(current_true)) + + loss = loss_raw * torch.tensor(class_weights) * torch.tensor(time_weights) + loss_list.append(loss.mean()) + + return torch.stack(loss_list).mean() + + return _loss_fn diff --git a/code/new_realtime/utils/model_utils.py b/code/new_realtime/utils/model_utils.py new file mode 100644 index 0000000..69cc9fc --- /dev/null +++ b/code/new_realtime/utils/model_utils.py @@ -0,0 +1,62 @@ +import os +import pickle +from datetime import datetime + +from utils.utils import get_config_id + + +def get_model_config(base_config: dict, + base_result_dir: str, + dataset_base_dir: str): + # get config identifier + # name_for_current_config = base_config["model_name"] + "_" + get_config_id(base_config) + date_part = datetime.now().strftime("%Y_%m_%d_%H_%M") + name_for_current_config = base_config["model_name"] + "_" + date_part + model_dir = os.path.abspath(f"{base_result_dir}/{name_for_current_config}") + if not os.path.exists(model_dir): + # create config + os.makedirs(model_dir) + + # append identifier to config + config = base_config.copy() + config["id"] = name_for_current_config + + # append paths + config["model_dir"] = model_dir + config["dataset_dir"] = os.path.join(dataset_base_dir, config["feature_config"]["feature_set_name"]) + config["feature_config"]["dataset_dir"] = config["dataset_dir"] + + # save model configuration + with open(f"{model_dir}/model_configuration.pickle", "wb") as f: + pickle.dump(config, f) + else: + # fetch config + with open(f"{model_dir}/model_configuration.pickle", "rb") as f: + config = pickle.load(f) + + # append paths + config["model_dir"] = model_dir + config["dataset_dir"] = os.path.join(dataset_base_dir, config["feature_config"]["feature_set_name"]) + config["feature_config"]["dataset_dir"] = config["dataset_dir"] + return config + + +def get_model_config_from_file(model_dir: str, + model_base_dir: str, + lmdb_base_dir: str): + # fetch config + with open(f"{model_dir}/model_configuration.pickle", "rb") as f: + config = pickle.load(f) + + # update paths based on base directories + config["model_dir"] = os.path.abspath(f"{model_base_dir}/{config['id']}") + config["feature_config"]["dataset_dir"] = os.path.abspath( + f"{lmdb_base_dir}/{config['feature_config']['feature_set_name']}") + + return config + + +def save_model_config(model_dir: str, model_config: dict): + # save model configuration + with open(f"{model_dir}/model_configuration.pickle", "wb") as f: + pickle.dump(model_config, f) diff --git a/code/new_realtime/utils/smoothing.py b/code/new_realtime/utils/smoothing.py new file mode 100644 index 0000000..ae02291 --- /dev/null +++ b/code/new_realtime/utils/smoothing.py @@ -0,0 +1,50 @@ +import numpy as np +from scipy.signal import butter, filtfilt +from statsmodels.tsa.stl._stl import STL + + +def highpass_filter(data, cutoff_freq, fs=288): + nyquist = 0.5 * fs + normal_cutoff = cutoff_freq / nyquist + b, a = butter(N=3, Wn=normal_cutoff, btype="high", analog=False) + return filtfilt(b, a, data) + + +def mirror_extend(series, extend_len): + """Mirrors the beginning and end of the time series to stabilize smoothing.""" + # Mirror extension + start_extension = series[:extend_len][::-1] # Reverse first part + end_extension = series[-extend_len:][::-1] # Reverse last part + + extended_series = np.concatenate([start_extension, series, end_extension]) + return extended_series + + +def get_trend(input_curve: np.ndarray | list, measurements_per_day: int = 288) -> np.ndarray: + extension_len = 3 + extended_input_curve = mirror_extend(input_curve, extension_len * measurements_per_day) + stl = STL(extended_input_curve, period=measurements_per_day, robust=False, trend=measurements_per_day * 14 + 1) + trend = stl.fit().trend + return trend[extension_len * measurements_per_day:-extension_len * measurements_per_day] + + +def get_curve_composition(input_curve: np.ndarray | list, measurements_per_day: int = 288) -> tuple: + """ + Decomposes the input curve into trend, seasonal, residual and smoothed components. + :param input_curve: raw input curve + :param measurements_per_day: seasonal period, here: measurements per day -> 288 + :return: composition of curve as tuple (trend, seasonal, residual, smoothed) + """ + extension_len = 3 + extended_input_curve = mirror_extend(input_curve, extension_len * measurements_per_day) + stl_results = STL(extended_input_curve, period=measurements_per_day, robust=False).fit() + long_term = (extended_input_curve - stl_results.seasonal) + wiggles = highpass_filter(long_term, 0.1) + smoothed = long_term - wiggles + + return ( + stl_results.trend[extension_len * measurements_per_day:-extension_len * measurements_per_day], + stl_results.seasonal[extension_len * measurements_per_day:-extension_len * measurements_per_day], + stl_results.resid[extension_len * measurements_per_day:-extension_len * measurements_per_day], + smoothed[extension_len * measurements_per_day:-extension_len * measurements_per_day], + ) diff --git a/code/new_realtime/utils/training.py b/code/new_realtime/utils/training.py new file mode 100644 index 0000000..20ff69d --- /dev/null +++ b/code/new_realtime/utils/training.py @@ -0,0 +1,330 @@ +import math + +import torch +from torch import nn +from torch.optim import AdamW +from torch.optim.lr_scheduler import OneCycleLR +from torch.utils.data import IterableDataset, DataLoader +from torch.utils.tensorboard import SummaryWriter +from tqdm import tqdm + +from utils.data_utils import LMDBIterableDataset +from utils.utils import get_logger + + +def process_tft_batch(model: nn.Module, + data_iterator: IterableDataset, + loss_functions: list, + device: torch.device, + model_configuration: dict) -> torch.Tensor: + batch = next(data_iterator) + batch = {k: v.to(device) for k, v in batch.items() if v is not None} + input_window_length = model_configuration["model_parameters"]["encoder_length"] + + preds = model(batch).cpu() # [B, decoder_len, Q] + target = batch["target"][:, input_window_length:, :].cpu() # match decoder segment + loss = get_x_y_loss(preds, target, loss_functions) + return loss + + +def get_x_y_loss(pred: torch.Tensor, + target: torch.Tensor, + loss_functions: list, + *args, **kwargs) -> torch.Tensor: + if len(loss_functions) > 1: + losses = list() + for dim in range(target.shape[-1]): + if len(pred.shape) > 2: + current_preds = pred[:, :, dim].ravel() + else: + current_preds = pred[:, dim] + if len(target.shape) > 2: + current_target = target[:, :, dim].ravel() + else: + current_target = target[:, dim].ravel() + # skip dimension, if it contains only NaN values, as loss cens + nan_indices = torch.isnan(current_target) + if torch.all(nan_indices): + continue + current_target = current_target[~nan_indices] + current_preds = current_preds[~nan_indices] + if len(current_target) == 0: + continue + + loss = loss_functions[dim](current_preds, current_target) + losses.append(loss) + loss = torch.stack(losses).mean() + else: + nan_indices = torch.isnan(target) + if torch.all(nan_indices): + return torch.tensor(0.0) + current_target = target[~nan_indices] + current_preds = pred[~nan_indices] + if len(current_target) == 0: + return torch.tensor(0.0) + loss = loss_functions[0](current_preds, current_target) + + return loss + + +def get_model_loss(model: nn.Module, + data_iterator: IterableDataset, + loss_functions: list, + device: str, + *args, **kwargs) -> torch.Tensor: + batch_x, batch_y = next(data_iterator) + batch_x = batch_x.to(device).float() + target = batch_y.to(device).float() + pred = model(batch_x) + loss = get_x_y_loss(pred, target, loss_functions) + return loss + + +def get_ranked_ids(all_ids, epoch, rank, world_size, base_seed=42): + """ + Get ranked ids for distributed training. + Args: + all_ids: list of all available ids + epoch: current epoch + rank: rank of the current process + world_size: number of processes + base_seed: base seed for random number generator + Returns: + list of ids for the current process + """ + g = torch.Generator() + g.manual_seed(base_seed + epoch) + permuted = torch.randperm(len(all_ids), generator=g).tolist() + return [all_ids[i] for i in permuted[rank::world_size]] + + +def train_model(model: nn.Module, + model_configuration: dict, + training_configuration: dict, + train_dataset: LMDBIterableDataset, + val_dataset: LMDBIterableDataset, + log_dir: str = "./logs", + logger=None) -> torch.nn.Module: + if logger is None: + logger = get_logger(__name__, f"{log_dir}/{model_configuration['id']}_{training_configuration['id']}.log") + + learning_parameters = training_configuration["learning_parameters"] + num_epochs = learning_parameters["epochs"] + patience = learning_parameters["patience"] + + training_id = training_configuration["id"] + + # get computation rank + if torch.distributed.is_initialized(): + local_rank = torch.distributed.get_rank() + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + world_size = torch.distributed.get_world_size() + else: + local_rank = 0 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + world_size = 1 + + logger.info(f"Rank {local_rank}: Using device: {device}, world size: {world_size}") + + # get subsets for distributed training + if torch.distributed.is_initialized(): + all_train_ids = train_dataset.lmdb_keys + train_subsets = [get_ranked_ids(all_train_ids, i, local_rank, world_size) for i in range(num_epochs)] + # calc total number of steps for gpu, as it is dependent on subsets + total_train_steps = sum([train_dataset.get_length_of_data_subset(subset) for subset in train_subsets]) + + all_val_ids = val_dataset.lmdb_keys + val_subsets = [get_ranked_ids(all_val_ids, i, local_rank, world_size) for i in range(num_epochs)] + else: + train_subsets = [train_dataset.lmdb_keys] * num_epochs + val_subsets = [val_dataset.lmdb_keys] * num_epochs + total_train_steps = len(train_dataset) + + # log the number of training steps for each epoch + train_subset_lengths = {f"epoch_{i}": len(subset) for i, subset in enumerate(train_subsets)} + logger.info(f"Train_subsets: {train_subset_lengths}") + logger.info(f"Rank {local_rank}: Total training steps: {total_train_steps}") + + # initialize loaders for non distributed + if not torch.distributed.is_initialized(): + # set the subsets for the datasets + train_dataset.set_key_subset(train_subsets[0]) + val_dataset.set_key_subset(val_subsets[0]) + + # create data loaders + train_dataloader = DataLoader( + train_dataset, + batch_size=None, + num_workers=4, + ) + val_dataloader = DataLoader( + val_dataset, + batch_size=None, + num_workers=4, + ) + + logger.info(f"Rank {local_rank}: Training {training_id} with {num_epochs} epochs") + logger.info(f"Rank {local_rank}: Training on {torch.cuda.device_count()} GPUs") + logger.info( + f"Rank {local_rank}: Current device: {torch.cuda.get_device_name(local_rank)} on local rank {local_rank}") + + model.to(device) + + # load training state from training configuration, if available + optimizer = AdamW(model.parameters(), lr=learning_parameters["learning_rate"]) + scheduler = OneCycleLR(optimizer, + max_lr=learning_parameters["learning_rate"], + # make sure to use length of full dataset here + total_steps=total_train_steps) + current_epoch = 1 + + loss_functions = training_configuration["loss_functions"] + # loss_fn = nn.MSELoss() + writer = SummaryWriter(log_dir=f'{log_dir}/{model_configuration["id"]}_{training_id}', ) + + best_val_loss = math.inf + epochs_no_improve = 0 + log_every_n_steps = max(len(train_dataset) // 500, 1) + + batch_loss_fn = model_configuration["batch_loss_fn"] + + for epoch in range(current_epoch, num_epochs + 1): + logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs}") + model.train() + total_train_loss = 0 + + # on distributed training, reshuffle the data + if torch.distributed.is_initialized(): + # update the datasets with the new ids + train_dataset.set_key_subset(train_subsets[epoch - 1]) + val_dataset.set_key_subset(val_subsets[epoch - 1]) + + # recreate data loaders + train_dataloader = DataLoader( + train_dataset, + batch_size=None, + num_workers=4, + ) + val_dataloader = DataLoader( + val_dataset, + batch_size=None, + num_workers=4, + ) + + iterator = iter(train_dataloader) + for step in tqdm(range(len(train_dataloader)), total=len(train_dataloader)): + loss = batch_loss_fn(model, + iterator, + loss_functions, + device, + model_configuration) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + scheduler.step() + + total_train_loss += loss.item() + if local_rank == 0: + if step % log_every_n_steps == 0: + writer.add_scalar("Loss/Train_Step", loss.item(), + ((epoch - 1) * len(train_dataloader) + step) * training_configuration[ + "batch_size"]) + writer.add_scalar("LR", scheduler.get_last_lr()[0], + ((epoch - 1) * len(train_dataloader) + step) * training_configuration[ + "batch_size"]) + writer.flush() + + avg_train_loss = total_train_loss / len(train_dataloader) + logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs} done. Train loss: {avg_train_loss:.4f}") + + # Validation + model.eval() + total_val_loss = 0 + logger.info(f"Rank {local_rank}: Validation") + with torch.no_grad(): + val_iter = iter(val_dataloader) + for step in tqdm(range(len(val_dataloader))): + loss = batch_loss_fn(model, + val_iter, + loss_functions, + device, + model_configuration) + + total_val_loss += loss.item() + + if len(val_dataloader) == 0: + logger.info(f"Rank {local_rank}: Validation set is empty, using 0 as validation loss.") + avg_val_loss = None + else: + avg_val_loss = total_val_loss / len(val_dataloader) + + # publish validation loss and wait for other gpus + if torch.distributed.is_initialized(): + if avg_val_loss is not None: + avg_val_loss_global = torch.tensor(avg_val_loss, device=device, dtype=torch.float32) + torch.distributed.all_reduce(avg_val_loss_global) + avg_val_loss_global /= torch.distributed.get_world_size() + + avg_train_loss_global = torch.tensor(avg_train_loss).to(device) + torch.distributed.all_reduce(avg_train_loss_global) + avg_train_loss_global /= torch.distributed.get_world_size() + else: + avg_val_loss_global = torch.tensor(avg_val_loss) + avg_train_loss_global = torch.tensor(avg_train_loss) + + # only rank 0 checks for early stopping + if local_rank == 0: + logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs} done. Val loss: {avg_val_loss:.4f}") + writer.add_scalar("Loss/Train_Epoch", avg_train_loss_global, epoch) + writer.add_scalar("Loss/Val_Epoch", avg_val_loss_global, epoch) + writer.flush() + + # Early stopping + should_stop = False + if avg_val_loss_global < best_val_loss: + logger.info( + f"Rank {local_rank}: Validation loss improved from {best_val_loss:.4f} to {avg_val_loss_global:.4f}.") + best_val_loss = avg_val_loss_global + epochs_no_improve = 0 + # torch.save(model.state_dict(), os.path.join(model_configuration["id"], "model.pt")) + save_fn = model_configuration["model_save_fn"] + save_fn(model, training_configuration) + else: + epochs_no_improve += 1 + logger.info( + f"Rank {local_rank}: No improvement in validation loss, no-improve count: {epochs_no_improve}") + if epochs_no_improve >= patience: + logger.info("Early stopping triggered.") + # broadcast stop signal to all gpus + should_stop = True + else: + should_stop = None + + if torch.distributed.is_initialized(): + if local_rank == 0: + should_stop_tensor = torch.tensor([int(should_stop)], device=device) + else: + should_stop_tensor = torch.zeros(1, dtype=torch.uint8, device=device) # safe default + torch.distributed.broadcast(should_stop_tensor, src=0) + should_stop = bool(should_stop_tensor.item()) + + if should_stop: + logger.info(f"Rank {local_rank}: Stopping training.") + break + + # ensure sync between epochs + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + # clean up + del loss + del train_dataloader + del val_dataloader + del model + del optimizer + del scheduler + + # free up memory + torch.cuda.synchronize() diff --git a/code/new_realtime/utils/training_utils.py b/code/new_realtime/utils/training_utils.py new file mode 100644 index 0000000..d0ae73f --- /dev/null +++ b/code/new_realtime/utils/training_utils.py @@ -0,0 +1,261 @@ +import os +import pickle +import random +from datetime import datetime + +import numpy as np +import torch +from torch import nn +from torch.nn import init +from sklearn.model_selection import train_test_split +import lmdb +from bson import ObjectId +from tqdm import tqdm + +from utils.lmdb_utils import get_lmdb_keys +from utils.utils import get_config_id +from utils.data_utils import produce_window_batches + +from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection + + +def get_training_config(base_config: dict, model_config: dict): + try: + if base_config["model_class"] != model_config["model_class"]: + raise Exception("Model type differ in model config and training config") + except KeyError: + raise Exception("Model class missing in training or model config") + + # training_config_id = get_config_id(base_config) + # hash = training_config_id[-5:] + + timestamp = datetime.now().strftime("%Y%m%d-%H%M") + training_config_id = f"{timestamp}" + + training_dir = os.path.abspath(f"{model_config['model_dir']}/trainings/{training_config_id}") + + # append identifier to config + training_config = base_config.copy() + training_config["id"] = training_config_id + training_config["training_dir"] = training_dir + + if not os.path.isdir(training_dir): + os.makedirs(training_dir) + # save training configuration + with open(f"{training_dir}/training_configuration.pickle", "wb") as f: + pickle.dump(training_config, f) + else: + # fetch config + with open(f"{training_dir}/training_configuration.pickle", "rb") as f: + training_config = pickle.load(f) + + return training_config + + +def get_training_config_from_file(training_dir: str, + base_model_dir: str, + model_configuration: dict) -> dict: + # fetch config + with open(f"{training_dir}/training_configuration.pickle", "rb") as f: + training_config = pickle.load(f) + + # update paths based on base directories + training_config["training_dir"] = os.path.join(os.path.abspath(base_model_dir), + model_configuration["id"], + "trainings", + training_config["id"]) + + return training_config + + +def weight_init(m): + """ + Usage: + model = Model() + model.apply(weight_init) + """ + if isinstance(m, nn.Conv1d): + init.normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.Conv2d): + init.xavier_normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.Conv3d): + init.xavier_normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose1d): + init.normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose2d): + init.xavier_normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.ConvTranspose3d): + init.xavier_normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.BatchNorm1d): + init.normal_(m.weight.data, mean=1, std=0.02) + init.constant_(m.bias.data, 0) + elif isinstance(m, nn.BatchNorm2d): + init.normal_(m.weight.data, mean=1, std=0.02) + init.constant_(m.bias.data, 0) + elif isinstance(m, nn.BatchNorm3d): + init.normal_(m.weight.data, mean=1, std=0.02) + init.constant_(m.bias.data, 0) + elif isinstance(m, nn.Linear): + init.xavier_normal_(m.weight.data) + if m.bias is not None: + init.normal_(m.bias.data) + elif isinstance(m, nn.LSTM): + for param in m.parameters(): + if len(param.shape) >= 2: + init.orthogonal_(param.data) + else: + init.normal_(param.data) + elif isinstance(m, nn.LSTMCell): + for param in m.parameters(): + if len(param.shape) >= 2: + init.orthogonal_(param.data) + else: + init.normal_(param.data) + elif isinstance(m, nn.GRU): + for param in m.parameters(): + if len(param.shape) >= 2: + init.orthogonal_(param.data) + else: + init.normal_(param.data) + for names in m._all_weights: + for name in filter(lambda n: "bias" in n, names): + bias = getattr(m, name) + n = bias.size(0) + bias.data[:n // 3].fill_(-1.) + elif isinstance(m, nn.GRUCell): + for param in m.parameters(): + if len(param.shape) >= 2: + init.orthogonal_(param.data) + else: + init.normal_(param.data) + + +def collate(batch_items: list) -> dict: + batch = dict() + for key in batch_items[0].keys(): + if key in ["combination_id", "time_index"]: + continue + else: + if batch_items[0][key] is None: + batch[key] = None + else: + batch[key] = np.stack([item[key] for item in batch_items]) + + for key in batch.keys(): + if batch[key] is not None: + batch[key] = torch.tensor(batch[key], dtype=torch.float32) + + return batch + + +def get_splits_by_user(input_keys: list, train_size: float, val_size: float, test_size: float): + if train_size + val_size + test_size != 1: + raise ValueError("Train, val and test sizes must sum to 1") + + if len(input_keys) == 0: + raise ValueError("Input keys list is empty") + + items_by_use = dict() + for input_key in tqdm(input_keys): + try: + user_id = get_cycles_collection().find_one({"_id": ObjectId(input_key)})["user_id"] + if user_id not in items_by_use: + items_by_use[user_id] = [] + items_by_use[user_id].append(input_key) + except Exception: + print(f"Error getting user id for key {input_key}") + continue + + user_ids = list(items_by_use.keys()) + random.shuffle(user_ids) + train_users, temp_users = train_test_split(user_ids, train_size=train_size, test_size=test_size + val_size) + # compute relative test size, as it must be relative to the remaining users + relative_test_size = test_size / (1 - train_size) + val_users, test_users = train_test_split(temp_users, test_size=relative_test_size) + + train_keys = [] + val_keys = [] + test_keys = [] + for user_id in train_users: + train_keys.extend(items_by_use[user_id]) + for user_id in val_users: + val_keys.extend(items_by_use[user_id]) + for user_id in test_users: + test_keys.extend(items_by_use[user_id]) + + return train_keys, val_keys, test_keys + + +def save_splits(train_keys: list, val_keys: list, test_keys: list, base_dir: str): + if not os.path.exists(base_dir): + os.makedirs(base_dir) + + with open(f"{base_dir}/train_keys.pickle", "wb") as f: + pickle.dump(train_keys, f) + with open(f"{base_dir}/val_keys.pickle", "wb") as f: + pickle.dump(val_keys, f) + with open(f"{base_dir}/test_keys.pickle", "wb") as f: + pickle.dump(test_keys, f) + + +def load_splits(base_dir: str): + with open(f"{base_dir}/train_keys.pickle", "rb") as f: + train_keys = pickle.load(f) + with open(f"{base_dir}/val_keys.pickle", "rb") as f: + val_keys = pickle.load(f) + with open(f"{base_dir}/test_keys.pickle", "rb") as f: + test_keys = pickle.load(f) + + return train_keys, val_keys, test_keys + + +def get_data_ids(model_configuration: dict, + training_configuration: dict, + env_path: str, + limit: int = None) -> tuple: + env = lmdb.open(f"{env_path}", readonly=True) + if os.path.exists(f"{model_configuration['feature_config']['dataset_dir']}/train_keys.pickle"): + train_ids, val_ids, test_ids = load_splits(model_configuration["feature_config"]["dataset_dir"]) + else: + lmdb_keys = get_lmdb_keys(env, limit) + # train_ids, val_ids, test_ids = get_splits_by_user(lmdb_keys, + # training_configuration["train_size"], + # training_configuration["val_size"], + # training_configuration["test_size"]) + train_ids, temp_ids = train_test_split(lmdb_keys, + train_size=training_configuration["train_size"], + test_size=training_configuration["val_size"] + training_configuration[ + "test_size"]) + # compute relative test size, as it must be relative to the remaining users + relative_test_size = training_configuration["test_size"] / (1 - training_configuration["train_size"]) + val_ids, test_ids = train_test_split(temp_ids, + test_size=relative_test_size) + # save splits to file + save_splits(train_ids, val_ids, test_ids, model_configuration["feature_config"]["dataset_dir"]) + + if limit is not None: + train_size = training_configuration["train_size"] + val_size = training_configuration["val_size"] + test_size = training_configuration["test_size"] + + train_lim = int(limit * train_size) + val_lim = int(limit * val_size) + test_lim = int(limit * test_size) + + train_ids = train_ids[:train_lim] + val_ids = val_ids[:val_lim] + test_ids = test_ids[:test_lim] + + return train_ids, val_ids, test_ids diff --git a/code/new_realtime/utils/utils.py b/code/new_realtime/utils/utils.py new file mode 100644 index 0000000..454880e --- /dev/null +++ b/code/new_realtime/utils/utils.py @@ -0,0 +1,152 @@ +import sys +import inspect +import hashlib +import logging +from functools import partial +import random +from typing import Callable +import importlib + +import numpy as np + + +def get_logger(module_name: str, filename: str = "main.log") -> logging.Logger: + """ + Returns a logger for the given module name and filename. + :param module_name: name of the module, as string + :param filename: name of the logging file, as string + :return: the logger, as logging.Logger object + """ + logger = logging.getLogger(module_name) + logger.setLevel(logging.DEBUG) + logger.propagate = False + formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + + file_handler = logging.FileHandler(filename) + file_handler.setLevel(logging.DEBUG) + file_handler.setFormatter(formatter) + + logger.addHandler(file_handler) + + # add system out handler + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setLevel(logging.INFO) + stream_handler.setFormatter(formatter) + + logger.addHandler(stream_handler) + + return logger + + +def get_variable_from_module(module_path: str, variable_name: str): + """ + Get a variable from a module by its name during runtime + + Args: + module_path: module to fetch variable from + variable_name: variable to fetch from module + + Returns: + variable from module + + """ + module = importlib.import_module(module_path) + variable = getattr(module, variable_name) + if variable is None: + raise ValueError(f"Variable {variable_name} not found in module {module_path}") + return variable + + +def get_callable_name(callable_obj): + """ + Get the name of the callable, handling `functools.partial`. + + :param callable_obj: callable object + :return: name of the callable + """ + if isinstance(callable_obj, partial): + func_name = callable_obj.func.__name__ + args = ", ".join(object_to_string(arg) for arg in callable_obj.args) + kwargs = ", ".join(f"{object_to_string(k)}={object_to_string(v)!r}" for k, v in callable_obj.keywords.items()) + return f"partial({func_name}, {args}, {kwargs})" + else: + if hasattr(callable_obj, '__name__'): + return callable_obj.__name__ + elif hasattr(callable_obj, '__class__'): + return callable_obj.__class__.__name__ + elif hasattr(callable_obj, '__hash__'): + return callable_obj.__hash__ + else: + raise ValueError(f"Could not determine name of callable object {callable_obj}") + + +def object_to_string(value, skip_types=None): + """ + Convert any object to a string representation that avoids memory addresses. + Handles complex data types recursively. + + :param value: object to convert + :param skip_types: types to skip during conversion + :return: string representation of the object + """ + if skip_types is None: + skip_types = [] + + if any(isinstance(value, t) for t in skip_types): + return 'skipped_type' + elif isinstance(value, (str, int, float, bool)): # Handle primitive data types directly + return repr(value) + elif isinstance(value, dict): + return '{' + ', '.join(f"{k}: {object_to_string(v, skip_types)}" for k, v in value.items()) + '}' + elif isinstance(value, (list, tuple)): + return '[' + ', '.join(object_to_string(item, skip_types) for item in value) + ']' + elif isinstance(value, partial): + return get_callable_name(value) + elif inspect.isclass(value): + return f"" + elif hasattr(value, + '__class__') and not value.__class__ != "function": # Correct handling for instances of classes, but not functions + return f"" + elif isinstance(value, Callable): + return f"" + else: + return repr(value) + + +def get_config_id(configuration: dict) -> str: + """ + Generate a somewhat unique human-readable model name from the model and training parameters. + :param configuration: dictionary containing model and training parameters + :return: human-readable model name + """ + adjectives = ["autumn", "hidden", "bitter", "misty", "silent", "empty", "dry", "dark", "summer", "icy", "delicate", + "quiet", "white", "black", "blue", "green", "red", "yellow", + "purple", "orange", "pink", "golden", "silver", "crimson", "violet", "azure", "amber", "sapphire", + "emerald", "ruby", "pearl", "topaz", "onyx", "turquoise", "citrine", ] + nouns = ["waterfall", "river", "breeze", "moon", "rain", "wind", "sea", "morning", "snow", "lake", "sunset", "pine", + "shadow", "leaf", "dawn", "glitter", "forest", "cloud", "sky", "sun", "butterfly", + "flower", "bird", "mountain", "valley", "ocean", "star", "night", "dream", "whisper", "echo", "horizon", + "wave", "petal", "dew", "mist"] + + # also add dataset config, but skip functions + base_name = object_to_string(configuration) + + # hash long name + basename_hash = hashlib.md5(base_name.encode()).hexdigest() + + # select adjective and noun based on hash + random.seed(int(basename_hash, 16)) + model_name = f"{random.choice(adjectives)}_{random.choice(nouns)}_{basename_hash[:5]}" + + return model_name + + +def convert_for_json(obj): + if isinstance(obj, dict): + return {k: convert_for_json(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [convert_for_json(v) for v in obj] + elif isinstance(obj, np.generic): + return obj.item() + else: + return obj diff --git a/code/new_realtime/utils/visualization.py b/code/new_realtime/utils/visualization.py new file mode 100644 index 0000000..adaaf79 --- /dev/null +++ b/code/new_realtime/utils/visualization.py @@ -0,0 +1,61 @@ +import numpy as np +import torch +from torch import nn + +from utils.dataset_creation import inverse_scale_feature + + +def plot_prediction_windows(index: int, + training_configuration: dict, + item_x: np.ndarray, + item_y: np.ndarray, + preds: np.ndarray, + scalers: dict, + features_to_plot: list, + output_feature_names: list, + fig_widget): + window_features = item_x[index] + indices = np.arange(window_features.shape[0]) + actual = item_y[index].ravel() + predicted = preds[index] + + # clear previous traces + fig_widget.data = [] + + for feature in features_to_plot: + feature_index = feature["index"] + feature_name = feature["name"] + fig_widget.add_scatter( + x=indices, + y=window_features[:, feature_index], + mode="lines", + name=feature_name, + ) + + num_outputs = predicted.shape[-1] + scaled_preds = list() + for i in range(num_outputs): + if isinstance(training_configuration["loss_functions"][i], nn.BCEWithLogitsLoss): + output = torch.sigmoid(torch.tensor(predicted[i])) + else: + output = predicted[i] + + output = float(output) + + scaled_output = inverse_scale_feature(output, + output_feature_names[i], + scalers) + + scaled_preds.append(scaled_output) + + scaled_actuals = list() + for i in range(num_outputs): + output = float(actual[i].numpy()) + scaled_output = inverse_scale_feature(output, + output_feature_names[i], + scalers) + + scaled_actuals.append(scaled_output) + + # add actual and predicted values + print(f"actual: {scaled_actuals}, predicted: {scaled_preds}")