From 8958cd7d95f61161f4dcfa821d058633bff5b715 Mon Sep 17 00:00:00 2001 From: czzhangheng Date: Mon, 18 Aug 2025 23:58:19 +0800 Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0.gitignore=E4=BB=A5=E5=BF=BD?= =?UTF-8?q?=E7=95=A5.temp=5Frepo=E6=96=87=E4=BB=B6=EF=BC=8C=E6=9B=B4?= =?UTF-8?q?=E6=96=B0README.md=E4=BB=A5=E6=94=AF=E6=8C=81=E6=96=B0=E6=A8=A1?= =?UTF-8?q?=E5=9E=8BSTAWnet=EF=BC=8C=E4=BF=AE=E6=94=B9DCRNN=E5=92=8CSTGNRD?= =?UTF-8?q?E=E9=85=8D=E7=BD=AE=E6=96=87=E4=BB=B6=EF=BC=8C=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0STAWnet=E5=92=8CSTGNRDE=E6=A8=A1=E5=9E=8B=E5=8F=8A?= =?UTF-8?q?=E5=85=B6=E8=AE=AD=E7=BB=83=E5=99=A8=EF=BC=8C=E5=A2=9E=E5=BC=BA?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E9=80=89=E6=8B=A9=E5=99=A8=E5=92=8C=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E5=8A=A0=E8=BD=BD=E5=99=A8=E4=BB=A5=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E6=96=B0=E6=A8=A1=E5=9E=8B=EF=BC=8C=E6=9B=B4=E6=96=B0=E6=8D=9F?= =?UTF-8?q?=E5=A4=B1=E5=87=BD=E6=95=B0=E4=BB=A5=E9=80=82=E5=BA=94=E6=96=B0?= =?UTF-8?q?=E9=9C=80=E6=B1=82=EF=BC=8C=E6=B7=BB=E5=8A=A0=E8=AE=AD=E7=BB=83?= =?UTF-8?q?=E7=BB=9F=E8=AE=A1=E5=8A=9F=E8=83=BD=E4=BB=A5=E8=B7=9F=E8=B8=AA?= =?UTF-8?q?=E8=AE=AD=E7=BB=83=E8=BF=87=E7=A8=8B=E4=B8=AD=E7=9A=84=E6=80=A7?= =?UTF-8?q?=E8=83=BD=E6=8C=87=E6=A0=87=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + README.md | 2 +- config/DCRNN/PEMSD4.yaml | 22 +- config/STAWnet/PEMSD3.yaml | 59 +++ config/STAWnet/PEMSD4.yaml | 60 +++ config/STAWnet/PEMSD7.yaml | 59 +++ config/STAWnet/PEMSD8.yaml | 59 +++ config/STGNRDE/PEMSD3.yaml | 53 +++ config/STGNRDE/PEMSD4.yaml | 53 +++ config/STGNRDE/PEMSD7.yaml | 53 +++ config/STGNRDE/PEMSD8.yaml | 53 +++ dataloader/loader_selector.py | 2 + lib/loss_function.py | 5 +- model/DCRNN/dcrnn_cell.py | 19 +- model/DCRNN/dcrnn_model.py | 32 +- model/STAWnet/STAWnet.py | 253 +++++++++++ model/STGNCDE/BasicTrainer_cde.py | 19 + model/STGNRDE/BasicTrainer_cde.py | 296 +++++++++++++ model/STGNRDE/GRDE.py | 120 +++++ model/STGNRDE/Make_model.py | 36 ++ model/STGNRDE/PEMSD4_GCDE.conf | 45 ++ model/STGNRDE/Run_cde.py | 205 +++++++++ model/STGNRDE/torchcde/__init__.py | 9 + model/STGNRDE/torchcde/interpolation_base.py | 22 + model/STGNRDE/torchcde/interpolation_cubic.py | 346 +++++++++++++++ .../interpolation_hermite_cubic_bdiff.py | 44 ++ .../STGNRDE/torchcde/interpolation_linear.py | 225 ++++++++++ model/STGNRDE/torchcde/log_ode.py | 134 ++++++ model/STGNRDE/torchcde/misc.py | 166 +++++++ model/STGNRDE/torchcde/solver.py | 409 ++++++++++++++++++ model/STGNRDE/vector_fields.py | 341 +++++++++++++++ model/model_selector.py | 4 + trainer/DCRNN_Trainer.py | 27 +- trainer/E32Trainer.py | 17 + trainer/EXP_trainer.py | 19 +- trainer/PDG2SEQ_Trainer.py | 18 + trainer/STMLP_Trainer.py | 17 + trainer/Trainer.py | 7 + trainer/Trainer_old.py | 17 + trainer/cdeTrainer/cdetrainer.py | 17 + trainer/trainer_selector.py | 2 + 41 files changed, 3313 insertions(+), 34 deletions(-) create mode 100644 config/STAWnet/PEMSD3.yaml create mode 100644 config/STAWnet/PEMSD4.yaml create mode 100644 config/STAWnet/PEMSD7.yaml create mode 100644 config/STAWnet/PEMSD8.yaml create mode 100644 config/STGNRDE/PEMSD3.yaml create mode 100644 config/STGNRDE/PEMSD4.yaml create mode 100644 config/STGNRDE/PEMSD7.yaml create mode 100644 config/STGNRDE/PEMSD8.yaml create mode 100644 model/STAWnet/STAWnet.py create mode 100755 model/STGNRDE/BasicTrainer_cde.py create mode 100755 model/STGNRDE/GRDE.py create mode 100755 model/STGNRDE/Make_model.py create mode 100755 model/STGNRDE/PEMSD4_GCDE.conf create mode 100755 model/STGNRDE/Run_cde.py create mode 100644 model/STGNRDE/torchcde/__init__.py create mode 100644 model/STGNRDE/torchcde/interpolation_base.py create mode 100644 model/STGNRDE/torchcde/interpolation_cubic.py create mode 100644 model/STGNRDE/torchcde/interpolation_hermite_cubic_bdiff.py create mode 100644 model/STGNRDE/torchcde/interpolation_linear.py create mode 100644 model/STGNRDE/torchcde/log_ode.py create mode 100644 model/STGNRDE/torchcde/misc.py create mode 100644 model/STGNRDE/torchcde/solver.py create mode 100755 model/STGNRDE/vector_fields.py diff --git a/.gitignore b/.gitignore index d67c4d1..dc5b170 100755 --- a/.gitignore +++ b/.gitignore @@ -170,3 +170,4 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ .DS_Store +.temp_repo/ \ No newline at end of file diff --git a/README.md b/README.md index 1a4ea28..4c4d154 100755 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ pip install pyyaml tqdm statsmodels h5py kagglehub torch torchvision torchaudio python run.py --model {model_name} --dataset {dataset_name} --mode {train, test} --device {cuda:0} ``` -- model_name: 目前支持：DSANET、STGCN、DCRNN、 GWN(GraphWaveNet)、STSGCN、AGCRN、STFGNN、STGODE、STGNCDE、DDGCRN、TWDGCN +- model_name: 目前支持：DSANET、STGCN、DCRNN、 GWN(GraphWaveNet)、STSGCN、AGCRN、STFGNN、STGODE、STGNCDE、DDGCRN、TWDGCN、STAWnet - dataset_name目前支持：PEMSD3，PEMSD4、PEMSD7、PEMSD8 - mode：train为训练模型，test为测试模型。测试模型需要在pre-train文件中找到模型的pth存档。 - device: 支持'cpu'、'cuda:0'、‘cuda:1’ ... 取决于机器卡数 diff --git a/config/DCRNN/PEMSD4.yaml b/config/DCRNN/PEMSD4.yaml index 60ba0e5..8758786 100755 --- a/config/DCRNN/PEMSD4.yaml +++ b/config/DCRNN/PEMSD4.yaml @@ -6,7 +6,7 @@ data: test_ratio: 0.2 tod: False normalizer: std - column_wise: False + column_wise: True default_graph: True add_time_in_day: True add_day_in_week: True @@ -21,26 +21,26 @@ model: max_diffusion_step: 2 cl_decay_steps: 1000 filter_type: dual_random_walk - num_rnn_layers: 1 + num_rnn_layers: 2 rnn_units: 64 seq_len: 12 use_curriculum_learning: True train: - loss_func: mae + loss_func: mask_mae seed: 10 batch_size: 64 epochs: 300 - lr_init: 0.003 - weight_decay: 0 - lr_decay: False - lr_decay_rate: 0.3 - lr_decay_step: "5,20,40,70" + lr_init: 0.001 + weight_decay: 0.0001 + lr_decay: True + lr_decay_rate: 0.1 + lr_decay_step: "10,20,40,80" early_stop: True - early_stop_patience: 15 - grad_norm: False + early_stop_patience: 25 + grad_norm: True max_grad_norm: 5 - real_value: True + real_value: False test: mae_thresh: null diff --git a/config/STAWnet/PEMSD3.yaml b/config/STAWnet/PEMSD3.yaml new file mode 100644 index 0000000..d54087e --- /dev/null +++ b/config/STAWnet/PEMSD3.yaml @@ -0,0 +1,59 @@ +data: + num_nodes: 358 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: False + normalizer: std + column_wise: False + default_graph: True + add_time_in_day: True + add_day_in_week: True + steps_per_day: 288 + days_per_week: 7 + +model: + input_dim: 1 + output_dim: 1 + dropout: 0.3 + gat_bool: True + addaptadj: True + aptonly: False + noapt: False + in_dim: 2 + residual_channels: 32 + dilation_channels: 32 + skip_channels: 256 + end_channels: 512 + kernel_size: 2 + blocks: 4 + layers: 2 + emb_length: 16 + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + weight_decay: 0 + lr_decay: False + lr_decay_rate: 0.3 + lr_decay_step: "5,20,40,70" + early_stop: True + early_stop_patience: 15 + grad_norm: False + max_grad_norm: 5 + real_value: True + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: False + + + diff --git a/config/STAWnet/PEMSD4.yaml b/config/STAWnet/PEMSD4.yaml new file mode 100644 index 0000000..b3346b1 --- /dev/null +++ b/config/STAWnet/PEMSD4.yaml @@ -0,0 +1,60 @@ +data: + num_nodes: 307 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: False + normalizer: std + column_wise: False + default_graph: True + add_time_in_day: True + add_day_in_week: True + steps_per_day: 288 + days_per_week: 7 + +model: + input_dim: 1 + output_dim: 1 + # STAWnet specific + dropout: 0.3 + gat_bool: True + addaptadj: True + aptonly: False + noapt: False + in_dim: 2 # 2 -> value + time_in_day; we also add day_in_week but not used by STAW core + residual_channels: 32 + dilation_channels: 32 + skip_channels: 256 + end_channels: 512 + kernel_size: 2 + blocks: 4 + layers: 2 + emb_length: 16 + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + weight_decay: 0 + lr_decay: False + lr_decay_rate: 0.3 + lr_decay_step: "5,20,40,70" + early_stop: True + early_stop_patience: 15 + grad_norm: False + max_grad_norm: 5 + real_value: True + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: False + + + diff --git a/config/STAWnet/PEMSD7.yaml b/config/STAWnet/PEMSD7.yaml new file mode 100644 index 0000000..a314c04 --- /dev/null +++ b/config/STAWnet/PEMSD7.yaml @@ -0,0 +1,59 @@ +data: + num_nodes: 883 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: False + normalizer: std + column_wise: False + default_graph: True + add_time_in_day: True + add_day_in_week: True + steps_per_day: 288 + days_per_week: 7 + +model: + input_dim: 1 + output_dim: 1 + dropout: 0.3 + gat_bool: True + addaptadj: True + aptonly: False + noapt: False + in_dim: 2 + residual_channels: 32 + dilation_channels: 32 + skip_channels: 256 + end_channels: 512 + kernel_size: 2 + blocks: 4 + layers: 2 + emb_length: 16 + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + weight_decay: 0 + lr_decay: False + lr_decay_rate: 0.3 + lr_decay_step: "5,20,40,70" + early_stop: True + early_stop_patience: 15 + grad_norm: False + max_grad_norm: 5 + real_value: True + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: False + + + diff --git a/config/STAWnet/PEMSD8.yaml b/config/STAWnet/PEMSD8.yaml new file mode 100644 index 0000000..75beec1 --- /dev/null +++ b/config/STAWnet/PEMSD8.yaml @@ -0,0 +1,59 @@ +data: + num_nodes: 170 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: False + normalizer: std + column_wise: False + default_graph: True + add_time_in_day: True + add_day_in_week: True + steps_per_day: 288 + days_per_week: 7 + +model: + input_dim: 1 + output_dim: 1 + dropout: 0.3 + gat_bool: True + addaptadj: True + aptonly: False + noapt: False + in_dim: 2 + residual_channels: 32 + dilation_channels: 32 + skip_channels: 256 + end_channels: 512 + kernel_size: 2 + blocks: 4 + layers: 2 + emb_length: 16 + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + weight_decay: 0 + lr_decay: False + lr_decay_rate: 0.3 + lr_decay_step: "5,20,40,70" + early_stop: True + early_stop_patience: 15 + grad_norm: False + max_grad_norm: 5 + real_value: True + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: False + + + diff --git a/config/STGNRDE/PEMSD3.yaml b/config/STGNRDE/PEMSD3.yaml new file mode 100644 index 0000000..3d98b05 --- /dev/null +++ b/config/STGNRDE/PEMSD3.yaml @@ -0,0 +1,53 @@ +data: + num_nodes: 358 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: false + normalizer: std + column_wise: false + default_graph: true + steps_per_day: 288 + days_per_week: 7 + +model: + g_type: agc + input_dim: 1 + output_dim: 1 + embed_dim: 10 + hid_dim: 128 + hid_hid_dim: 128 + num_layers: 2 + cheb_k: 3 + solver: rk4 + model_type: rde + interpolation: cubic + emb_opt: false + adp_opt: false + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 100 + lr_init: 0.003 + lr_decay: false + weight_decay: 0 + lr_decay_rate: 0.3 + lr_decay_step: [5, 20, 40, 70] + early_stop: true + early_stop_patience: 15 + grad_norm: false + max_grad_norm: 5 + real_value: true + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: false + + diff --git a/config/STGNRDE/PEMSD4.yaml b/config/STGNRDE/PEMSD4.yaml new file mode 100644 index 0000000..c75617c --- /dev/null +++ b/config/STGNRDE/PEMSD4.yaml @@ -0,0 +1,53 @@ +data: + num_nodes: 307 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: false + normalizer: std + column_wise: false + default_graph: true + steps_per_day: 288 + days_per_week: 7 + +model: + g_type: agc + input_dim: 2 + output_dim: 1 + embed_dim: 10 + hid_dim: 128 + hid_hid_dim: 128 + num_layers: 2 + cheb_k: 3 + solver: rk4 + model_type: rde + interpolation: cubic + emb_opt: false + adp_opt: false + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 100 + lr_init: 0.003 + lr_decay: false + weight_decay: 0 + lr_decay_rate: 0.3 + lr_decay_step: [5, 20, 40, 70] + early_stop: true + early_stop_patience: 15 + grad_norm: false + max_grad_norm: 5 + real_value: true + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: false + + diff --git a/config/STGNRDE/PEMSD7.yaml b/config/STGNRDE/PEMSD7.yaml new file mode 100644 index 0000000..e135a63 --- /dev/null +++ b/config/STGNRDE/PEMSD7.yaml @@ -0,0 +1,53 @@ +data: + num_nodes: 883 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: false + normalizer: std + column_wise: false + default_graph: true + steps_per_day: 288 + days_per_week: 7 + +model: + g_type: agc + input_dim: 1 + output_dim: 1 + embed_dim: 10 + hid_dim: 128 + hid_hid_dim: 128 + num_layers: 2 + cheb_k: 3 + solver: rk4 + model_type: rde + interpolation: cubic + emb_opt: false + adp_opt: false + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + lr_decay: false + weight_decay: 0 + lr_decay_rate: 0.3 + lr_decay_step: [5, 20, 40, 70] + early_stop: true + early_stop_patience: 15 + grad_norm: false + max_grad_norm: 5 + real_value: true + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: false + + diff --git a/config/STGNRDE/PEMSD8.yaml b/config/STGNRDE/PEMSD8.yaml new file mode 100644 index 0000000..3c0751b --- /dev/null +++ b/config/STGNRDE/PEMSD8.yaml @@ -0,0 +1,53 @@ +data: + num_nodes: 170 + lag: 12 + horizon: 12 + val_ratio: 0.2 + test_ratio: 0.2 + tod: false + normalizer: std + column_wise: false + default_graph: true + steps_per_day: 288 + days_per_week: 7 + +model: + g_type: agc + input_dim: 2 + output_dim: 1 + embed_dim: 10 + hid_dim: 128 + hid_hid_dim: 128 + num_layers: 2 + cheb_k: 3 + solver: rk4 + model_type: rde + interpolation: cubic + emb_opt: false + adp_opt: false + +train: + loss_func: mae + seed: 10 + batch_size: 64 + epochs: 300 + lr_init: 0.003 + lr_decay: false + weight_decay: 0 + lr_decay_rate: 0.3 + lr_decay_step: [5, 20, 40, 70] + early_stop: true + early_stop_patience: 15 + grad_norm: false + max_grad_norm: 5 + real_value: true + +test: + mae_thresh: null + mape_thresh: 0.0 + +log: + log_step: 200 + plot: false + + diff --git a/dataloader/loader_selector.py b/dataloader/loader_selector.py index eea8d60..b7f697c 100755 --- a/dataloader/loader_selector.py +++ b/dataloader/loader_selector.py @@ -2,10 +2,12 @@ from dataloader.cde_loader.cdeDataloader import get_dataloader as cde_loader from dataloader.PeMSDdataloader import get_dataloader as normal_loader from dataloader.DCRNNdataloader import get_dataloader as DCRNN_loader from dataloader.EXPdataloader import get_dataloader as EXP_loader +from dataloader.cde_loader.cdeDataloader import get_dataloader as nrde_loader def get_dataloader(config, normalizer, single): match config['model']['type']: case 'STGNCDE': return cde_loader(config['data'], normalizer, single) + case 'STGNRDE': return nrde_loader(config['data'], normalizer, single) case 'DCRNN': return DCRNN_loader(config['data'], normalizer, single) case 'EXP': return EXP_loader(config['data'], normalizer, single) case _: return normal_loader(config['data'], normalizer, single) diff --git a/lib/loss_function.py b/lib/loss_function.py index 6645e11..ad8caf4 100755 --- a/lib/loss_function.py +++ b/lib/loss_function.py @@ -1,9 +1,9 @@ def masked_mae_loss(scaler, mask_value): def loss(preds, labels): + # 仅对预测反归一化；标签在数据管道中保持原始量纲 if scaler: preds = scaler.inverse_transform(preds) - labels = scaler.inverse_transform(labels) return mae_torch(pred=preds, true=labels, mask_value=mask_value) return loss @@ -11,7 +11,8 @@ def masked_mae_loss(scaler, mask_value): def get_loss_function(args, scaler): if args['loss_func'] == 'mask_mae': - return masked_mae_loss(scaler, mask_value=0.0).to(args['device']) + # Return callable loss (no .to for function closures); disable masking by default + return masked_mae_loss(scaler, mask_value=None) elif args['loss_func'] == 'mae': return torch.nn.L1Loss().to(args['device']) elif args['loss_func'] == 'mse': diff --git a/model/DCRNN/dcrnn_cell.py b/model/DCRNN/dcrnn_cell.py index c425069..fabe400 100755 --- a/model/DCRNN/dcrnn_cell.py +++ b/model/DCRNN/dcrnn_cell.py @@ -34,7 +34,7 @@ class LayerParams: class DCGRUCell(torch.nn.Module): - def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, nonlinearity='tanh', + def __init__(self, num_units, adj_mx, max_diffusion_step, num_nodes, input_dim=None, nonlinearity='tanh', filter_type="laplacian", use_gc_for_ru=True): """ @@ -55,6 +55,7 @@ class DCGRUCell(torch.nn.Module): self._max_diffusion_step = max_diffusion_step self._supports = [] self._use_gc_for_ru = use_gc_for_ru + self._input_dim = input_dim # optional; if None, will be inferred at first forward supports = [] if filter_type == "laplacian": supports.append(utils.calculate_scaled_laplacian(adj_mx, lambda_max=None)) @@ -71,6 +72,22 @@ class DCGRUCell(torch.nn.Module): self._fc_params = LayerParams(self, 'fc') self._gconv_params = LayerParams(self, 'gconv') + # Pre-create parameters if input_dim is known + if self._input_dim is not None: + num_matrices = len(self._supports) * self._max_diffusion_step + 1 + input_size = self._input_dim + self._num_units + # FC weights/biases for RU gates (2 * num_units) + self._fc_params.get_weights((input_size, 2 * self._num_units)) + self._fc_params.get_biases(2 * self._num_units, bias_start=1.0) + # Optionally for candidate (num_units) if FC path is used + self._fc_params.get_weights((input_size, self._num_units)) + self._fc_params.get_biases(self._num_units, bias_start=0.0) + # GConv weights/biases for RU gates and candidate + self._gconv_params.get_weights((input_size * num_matrices, 2 * self._num_units)) + self._gconv_params.get_biases(2 * self._num_units, bias_start=1.0) + self._gconv_params.get_weights((input_size * num_matrices, self._num_units)) + self._gconv_params.get_biases(self._num_units, bias_start=0.0) + @staticmethod def _build_sparse_matrix(L): L = L.tocoo() diff --git a/model/DCRNN/dcrnn_model.py b/model/DCRNN/dcrnn_model.py index eb5803a..f7e7648 100755 --- a/model/DCRNN/dcrnn_model.py +++ b/model/DCRNN/dcrnn_model.py @@ -15,6 +15,10 @@ class Seq2SeqAttrs: self.num_nodes = args.get('num_nodes', 1) self.num_rnn_layers = args.get('num_rnn_layers', 1) self.rnn_units = args.get('rnn_units') + self.input_dim = args.get('input_dim', 1) + self.output_dim = args.get('output_dim', 1) + self.horizon = args.get('horizon', 12) + self.seq_len = args.get('seq_len', 12) self.hidden_state_size = self.num_nodes * self.rnn_units @@ -26,7 +30,7 @@ class EncoderModel(nn.Module, Seq2SeqAttrs): self.seq_len = args.get('seq_len') # for the encoder self.dcgru_layers = nn.ModuleList( [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes, - filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) + input_dim=self.input_dim, filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) def forward(self, inputs, hidden_state=None): """ @@ -63,7 +67,7 @@ class DecoderModel(nn.Module, Seq2SeqAttrs): self.projection_layer = nn.Linear(self.rnn_units, self.output_dim) self.dcgru_layers = nn.ModuleList( [DCGRUCell(self.rnn_units, adj_mx, self.max_diffusion_step, self.num_nodes, - filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) + input_dim=self.output_dim, filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) def forward(self, inputs, hidden_state=None): """ @@ -146,17 +150,19 @@ class DCRNNModel(nn.Module, Seq2SeqAttrs): def forward(self, inputs, labels=None): """ - seq2seq forward pass 64 12 307 3 - :param inputs: shape (seq_len, batch_size, num_sensor * input_dim) 12 64 307 * 1 - :param labels: shape (horizon, batch_size, num_sensor * output) 12 64 307 1 - :param batches_seen: batches seen till now - :return: output: (self.horizon, batch_size, self.num_nodes * self.output_dim) + seq2seq forward pass. inputs: [B, T, N, C] """ - inputs = inputs[..., 0].permute(1, 0, 2) - labels = labels[..., 0].permute(1, 0, 2) - encoder_hidden_state = self.encoder(inputs) - outputs = self.decoder(encoder_hidden_state, labels, batches_seen=self.batch_seen) + x = inputs[..., :self.input_dim] + x = x.permute(1, 0, 2, 3).contiguous().view(self.seq_len, -1, self.num_nodes * self.input_dim) + + y = None + if labels is not None: + y = labels[..., :self.output_dim] + y = y.permute(1, 0, 2, 3).contiguous().view(self.horizon, -1, self.num_nodes * self.output_dim) + + encoder_hidden_state = self.encoder(x) + outputs = self.decoder(encoder_hidden_state, y, batches_seen=self.batch_seen) self.batch_seen += 1 - outputs = outputs.unsqueeze(dim=-1) # [12,64,307,1] - outputs = outputs.permute(1, 0, 2, 3) # [64,12,307,1] + outputs = outputs.view(self.horizon, -1, self.num_nodes, self.output_dim) + outputs = outputs.permute(1, 0, 2, 3).contiguous() return outputs diff --git a/model/STAWnet/STAWnet.py b/model/STAWnet/STAWnet.py new file mode 100644 index 0000000..b36aef3 --- /dev/null +++ b/model/STAWnet/STAWnet.py @@ -0,0 +1,253 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class LinearConv2d(nn.Module): + def __init__(self, c_in, c_out): + super().__init__() + self.mlp = nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0, 0), stride=(1, 1), bias=True) + + def forward(self, x): + return self.mlp(x) + + +class PointwiseConv2d(nn.Module): + def __init__(self, c_in, c_out): + super().__init__() + self.mlp = nn.Conv2d(c_in, c_out, kernel_size=(1, 1), padding=(0, 0), stride=(1, 1), bias=True) + + def forward(self, x): + return self.mlp(x) + + +class GraphAttention(nn.Module): + def __init__(self, c_in, c_out, dropout, d=16, emb_length=0, aptonly=False, noapt=False): + super().__init__() + self.d = d + self.aptonly = aptonly + self.noapt = noapt + self.mlp = LinearConv2d(c_in * 2, c_out) + self.dropout = dropout + self.emb_length = emb_length + + if aptonly: + self.qm = PointwiseConv2d(self.emb_length, d) + self.km = PointwiseConv2d(self.emb_length, d) + elif noapt: + self.qm = PointwiseConv2d(c_in, d) + self.km = PointwiseConv2d(c_in, d) + else: + self.qm = PointwiseConv2d(c_in + self.emb_length, d) + self.km = PointwiseConv2d(c_in + self.emb_length, d) + + def forward(self, x, embedding): + # x: [B, C, N, T] + # embedding: [emb_length, N] (as parameter), we broadcast to [B, emb_length, N, T] + out = [x] + + embedding = embedding.repeat((x.shape[0], x.shape[-1], 1, 1)) # [B, T, emb, N] + embedding = embedding.permute(0, 2, 3, 1).contiguous() # [B, emb, N, T] + + if self.aptonly: + x_embedding = embedding + query = self.qm(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + key = self.km(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + attention = torch.matmul(query, key.permute(0, 1, 3, 2)) # [B, T, N, N] + attention = attention / (self.d ** 0.5) + attention = F.softmax(attention, dim=-1) + elif self.noapt: + x_embedding = x + query = self.qm(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + key = self.km(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + attention = torch.matmul(query, key.permute(0, 1, 3, 2)) # [B, T, N, N] + attention = attention / (self.d ** 0.5) + attention = F.softmax(attention, dim=-1) + else: + x_embedding = torch.cat([x, embedding], dim=1) # [B, C+emb, N, T] + query = self.qm(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + key = self.km(x_embedding).permute(0, 3, 2, 1) # [B, T, N, d] + attention = torch.matmul(query, key.permute(0, 1, 3, 2)) # [B, T, N, N] + attention = attention / (self.d ** 0.5) + attention = F.softmax(attention, dim=-1) + + # apply attention over nodes: [B, C, N, T] -> [B, T, C, N] * [B, T, N, N] -> [B, T, C, N] + x = torch.matmul(x.permute(0, 3, 1, 2), attention).permute(0, 2, 3, 1) + out.append(x) + + h = torch.cat(out, dim=1) + h = self.mlp(h) + h = F.dropout(h, self.dropout, training=self.training) + return h + + +class STAWnetCore(nn.Module): + def __init__( + self, + device, + num_nodes, + dropout=0.3, + gat_bool=True, + addaptadj=True, + aptonly=False, + noapt=False, + in_dim=2, + out_dim=12, + residual_channels=32, + dilation_channels=32, + skip_channels=256, + end_channels=512, + kernel_size=2, + blocks=4, + layers=2, + emb_length=16, + ): + super().__init__() + + self.dropout = dropout + self.blocks = blocks + self.layers = layers + self.gat_bool = gat_bool + self.aptonly = aptonly + self.noapt = noapt + self.addaptadj = addaptadj + self.emb_length = emb_length + + self.filter_convs = nn.ModuleList() + self.gate_convs = nn.ModuleList() + self.residual_convs = nn.ModuleList() + self.skip_convs = nn.ModuleList() + self.bn = nn.ModuleList() + self.gat = nn.ModuleList() + + self.start_conv = nn.Conv2d(in_channels=in_dim, out_channels=residual_channels, kernel_size=(1, 1)) + + self.supports = None + receptive_field = 1 + + if gat_bool and addaptadj: + # learnable node embeddings: [emb_length, N] + self.embedding = nn.Parameter(torch.randn(self.emb_length, num_nodes, device=device), requires_grad=True) + + for _ in range(blocks): + additional_scope = kernel_size - 1 + new_dilation = 1 + for _ in range(layers): + # dilated temporal convs + self.filter_convs.append( + nn.Conv2d(in_channels=residual_channels, out_channels=dilation_channels, kernel_size=(1, kernel_size), dilation=new_dilation) + ) + self.gate_convs.append( + nn.Conv2d(in_channels=residual_channels, out_channels=dilation_channels, kernel_size=(1, kernel_size), dilation=new_dilation) + ) + + # 1x1 residual/skip + self.residual_convs.append(nn.Conv2d(in_channels=dilation_channels, out_channels=residual_channels, kernel_size=(1, 1))) + self.skip_convs.append(nn.Conv2d(in_channels=dilation_channels, out_channels=skip_channels, kernel_size=(1, 1))) + + self.bn.append(nn.BatchNorm2d(residual_channels)) + + new_dilation *= 2 + receptive_field += additional_scope + additional_scope *= 2 + + if self.gat_bool: + self.gat.append(GraphAttention(dilation_channels, residual_channels, dropout, emb_length=emb_length, aptonly=aptonly, noapt=noapt)) + + self.end_conv_1 = nn.Conv2d(in_channels=skip_channels, out_channels=end_channels, kernel_size=(1, 1), bias=True) + self.end_conv_2 = nn.Conv2d(in_channels=end_channels, out_channels=out_dim, kernel_size=(1, 1), bias=True) + + self.receptive_field = receptive_field + + def forward(self, input): + # input: [B, C_in, N, T] + in_len = input.size(3) + if in_len < self.receptive_field: + x = F.pad(input, (self.receptive_field - in_len, 0, 0, 0)) + else: + x = input + + x = self.start_conv(x) + skip = 0 + + for i in range(self.blocks * self.layers): + residual = x + # gated temporal conv + filt = torch.tanh(self.filter_convs[i](residual)) + gate = torch.sigmoid(self.gate_convs[i](residual)) + x = filt * gate + + # skip connection accumulation (align time length) + s = self.skip_convs[i](x) + if isinstance(skip, torch.Tensor): + skip = skip[:, :, :, -s.size(3):] + else: + skip = 0 + skip = s + skip + + # spatial attention or residual conv + if self.gat_bool and hasattr(self, 'embedding'): + x = self.gat[i](x, self.embedding) + else: + x = self.residual_convs[i](x) + + # residual connection and BN + x = x + residual[:, :, :, -x.size(3):] + x = self.bn[i](x) + + x = F.relu(skip) + x = F.relu(self.end_conv_1(x)) + x = self.end_conv_2(x) + # shape: [B, horizon(out_dim), N, T_reduced(usually 1)] + return x + + +class STAWnet(nn.Module): + """ + Project-adapted STAWnet wrapper that matches the common interface: + - Input: [B, T, N, C_total] + - Output: [B, horizon, N, output_dim] + """ + + def __init__(self, args): + super().__init__() + self.args = args + + device = args.get('device', 'cpu') + num_nodes = args['num_nodes'] + + # Model IO configs + in_dim = args.get('in_dim', 2) # how many covariates to feed into the model + horizon = args.get('horizon', 12) + output_dim = args.get('output_dim', 1) + + self.use_channels = in_dim + + self.core = STAWnetCore( + device=device, + num_nodes=num_nodes, + dropout=args.get('dropout', 0.3), + gat_bool=args.get('gat_bool', True), + addaptadj=args.get('addaptadj', True), + aptonly=args.get('aptonly', False), + noapt=args.get('noapt', False), + in_dim=in_dim, + out_dim=horizon, # channels represent horizon steps + residual_channels=args.get('residual_channels', 32), + dilation_channels=args.get('dilation_channels', 32), + skip_channels=args.get('skip_channels', 256), + end_channels=args.get('end_channels', 512), + kernel_size=args.get('kernel_size', 2), + blocks=args.get('blocks', 4), + layers=args.get('layers', 2), + emb_length=args.get('emb_length', 16), + ) + + def forward(self, x): + # x: [B, T, N, C_total] -> pick first self.use_channels, then to [B, C, N, T] + x = x[..., :self.use_channels].transpose(1, 3) # [B, C, N, T] + y = self.core(x) # [B, horizon, N, T_reduced(=1)] + # Keep horizon on channel dimension to match project convention (like GWN) + return y + + diff --git a/model/STGNCDE/BasicTrainer_cde.py b/model/STGNCDE/BasicTrainer_cde.py index ef0c4f3..569bcd5 100755 --- a/model/STGNCDE/BasicTrainer_cde.py +++ b/model/STGNCDE/BasicTrainer_cde.py @@ -7,6 +7,7 @@ import numpy as np from lib.logger import get_logger from lib.metrics import All_Metrics from lib.TrainInits import print_model_parameters +from lib.training_stats import TrainingStats class Trainer(object): def __init__(self, model, vector_field_f, vector_field_g, loss, optimizer, train_loader, val_loader, test_loader, @@ -42,6 +43,8 @@ class Trainer(object): self.device = device self.times = times.to(self.device, dtype=torch.float) self.w = w + # Stats tracker + self.stats = TrainingStats(device=device) def val_epoch(self, epoch, val_dataloader): self.model.eval() @@ -49,6 +52,7 @@ class Trainer(object): with torch.no_grad(): for batch_idx, batch in enumerate(self.val_loader): + start_time = time.time() # for iter, batch in enumerate(val_dataloader): batch = tuple(b.to(self.device, dtype=torch.float) for b in batch) *valid_coeffs, target = batch @@ -61,8 +65,11 @@ class Trainer(object): #a whole batch of Metr_LA is filtered if not torch.isnan(loss): total_val_loss += loss.item() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, 'val') val_loss = total_val_loss / len(val_dataloader) self.logger.info('**********Val Epoch {}: average Loss: {:.6f}'.format(epoch, val_loss)) + self.stats.record_memory_usage() if self.args.tensorboard: self.w.add_scalar(f'valid/loss', val_loss, epoch) return val_loss @@ -73,6 +80,7 @@ class Trainer(object): # for batch_idx, (data, target) in enumerate(self.train_loader): # for batch_idx, (data, target) in enumerate(self.train_loader): for batch_idx, batch in enumerate(self.train_loader): + start_time = time.time() batch = tuple(b.to(self.device, dtype=torch.float) for b in batch) *train_coeffs, target = batch # data = data[..., :self.args.input_dim] @@ -101,6 +109,8 @@ class Trainer(object): if self.args.grad_norm: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, 'train') total_loss += loss.item() #log information @@ -109,6 +119,7 @@ class Trainer(object): epoch, batch_idx, self.train_per_epoch, loss.item())) train_epoch_loss = total_loss/self.train_per_epoch self.logger.info('**********Train Epoch {}: averaged Loss: {:.6f}'.format(epoch, train_epoch_loss)) + self.stats.record_memory_usage() if self.args.tensorboard: self.w.add_scalar(f'train/loss', train_epoch_loss, epoch) @@ -123,6 +134,7 @@ class Trainer(object): not_improved_count = 0 train_loss_list = [] val_loss_list = [] + self.stats.start_training() start_time = time.time() for epoch in range(1, self.args.epochs + 1): #epoch_time = time.time() @@ -168,6 +180,13 @@ class Trainer(object): training_time = time.time() - start_time self.logger.info("Total training time: {:.4f}min, best loss: {:.6f}".format((training_time / 60), best_loss)) + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass #save the best model to file if not self.args.debug: diff --git a/model/STGNRDE/BasicTrainer_cde.py b/model/STGNRDE/BasicTrainer_cde.py new file mode 100755 index 0000000..569bcd5 --- /dev/null +++ b/model/STGNRDE/BasicTrainer_cde.py @@ -0,0 +1,296 @@ +import torch +import math +import os +import time +import copy +import numpy as np +from lib.logger import get_logger +from lib.metrics import All_Metrics +from lib.TrainInits import print_model_parameters +from lib.training_stats import TrainingStats + +class Trainer(object): + def __init__(self, model, vector_field_f, vector_field_g, loss, optimizer, train_loader, val_loader, test_loader, + scaler, args, lr_scheduler, device, times, + w): + super(Trainer, self).__init__() + self.model = model + self.vector_field_f = vector_field_f + self.vector_field_g = vector_field_g + self.loss = loss + self.optimizer = optimizer + self.train_loader = train_loader + self.val_loader = val_loader + self.test_loader = test_loader + self.scaler = scaler + self.args = args + self.lr_scheduler = lr_scheduler + self.train_per_epoch = len(train_loader) + if val_loader != None: + self.val_per_epoch = len(val_loader) + self.best_path = os.path.join(self.args.log_dir, 'best_model.pth') + self.loss_figure_path = os.path.join(self.args.log_dir, 'loss.png') + #log + if os.path.isdir(args.log_dir) == False and not args.debug: + os.makedirs(args.log_dir, exist_ok=True) + self.logger = get_logger(args.log_dir, name=args.model, debug=args.debug) + self.logger.info('Experiment log path in: {}'.format(args.log_dir)) + total_param = print_model_parameters(model, only_num=False) + for arg, value in sorted(vars(args).items()): + self.logger.info("Argument %s: %r", arg, value) + self.logger.info(self.model) + self.logger.info("Total params: {}".format(str(total_param))) + self.device = device + self.times = times.to(self.device, dtype=torch.float) + self.w = w + # Stats tracker + self.stats = TrainingStats(device=device) + + def val_epoch(self, epoch, val_dataloader): + self.model.eval() + total_val_loss = 0 + + with torch.no_grad(): + for batch_idx, batch in enumerate(self.val_loader): + start_time = time.time() + # for iter, batch in enumerate(val_dataloader): + batch = tuple(b.to(self.device, dtype=torch.float) for b in batch) + *valid_coeffs, target = batch + # data = data[..., :self.args.input_dim] + label = target[..., :self.args.output_dim] + output = self.model(self.times, valid_coeffs) + if self.args.real_value: + label = self.scaler.inverse_transform(label) + loss = self.loss(output.cuda(), label) + #a whole batch of Metr_LA is filtered + if not torch.isnan(loss): + total_val_loss += loss.item() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, 'val') + val_loss = total_val_loss / len(val_dataloader) + self.logger.info('**********Val Epoch {}: average Loss: {:.6f}'.format(epoch, val_loss)) + self.stats.record_memory_usage() + if self.args.tensorboard: + self.w.add_scalar(f'valid/loss', val_loss, epoch) + return val_loss + + def train_epoch(self, epoch): + self.model.train() + total_loss = 0 + # for batch_idx, (data, target) in enumerate(self.train_loader): + # for batch_idx, (data, target) in enumerate(self.train_loader): + for batch_idx, batch in enumerate(self.train_loader): + start_time = time.time() + batch = tuple(b.to(self.device, dtype=torch.float) for b in batch) + *train_coeffs, target = batch + # data = data[..., :self.args.input_dim] + label = target[..., :self.args.output_dim] # (..., 1) + self.optimizer.zero_grad() + + # #teacher_forcing for RNN encoder-decoder model + # #if teacher_forcing_ratio = 1: use label as input in the decoder for all steps + # if self.args.teacher_forcing: + # global_step = (epoch - 1) * self.train_per_epoch + batch_idx + # teacher_forcing_ratio = self._compute_sampling_threshold(global_step, self.args.tf_decay_steps) + # else: + # teacher_forcing_ratio = 1. + #data and target shape: B, T, N, F; output shape: B, T, N, F + output = self.model(self.times, train_coeffs) + # output = self.model(train_coeffs, target, teacher_forcing_ratio=teacher_forcing_ratio) + if self.args.real_value: + label = self.scaler.inverse_transform(label) + loss = self.loss(output.cuda(), label) + + # loss = _add_weight_regularisation(loss, self.vector_field_g) #TODO: regularization + # loss = _add_weight_regularisation(loss, self.vector_field_f) #TODO: regularization + loss.backward() + + # add max grad clipping + if self.args.grad_norm: + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm) + self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, 'train') + total_loss += loss.item() + + #log information + if batch_idx % self.args.log_step == 0: + self.logger.info('Train Epoch {}: {}/{} Loss: {:.6f}'.format( + epoch, batch_idx, self.train_per_epoch, loss.item())) + train_epoch_loss = total_loss/self.train_per_epoch + self.logger.info('**********Train Epoch {}: averaged Loss: {:.6f}'.format(epoch, train_epoch_loss)) + self.stats.record_memory_usage() + if self.args.tensorboard: + self.w.add_scalar(f'train/loss', train_epoch_loss, epoch) + + #learning rate decay + if self.args.lr_decay: + self.lr_scheduler.step() + return train_epoch_loss + + def train(self): + best_model = None + best_loss = float('inf') + not_improved_count = 0 + train_loss_list = [] + val_loss_list = [] + self.stats.start_training() + start_time = time.time() + for epoch in range(1, self.args.epochs + 1): + #epoch_time = time.time() + train_epoch_loss = self.train_epoch(epoch) + #print(time.time()-epoch_time) + #exit() + if self.val_loader == None: + val_dataloader = self.test_loader + else: + val_dataloader = self.val_loader + val_epoch_loss = self.val_epoch(epoch, val_dataloader) + + #print('LR:', self.optimizer.param_groups[0]['lr']) + train_loss_list.append(train_epoch_loss) + val_loss_list.append(val_epoch_loss) + if train_epoch_loss > 1e6: + self.logger.warning('Gradient explosion detected. Ending...') + break + #if self.val_loader == None: + #val_epoch_loss = train_epoch_loss + if val_epoch_loss < best_loss: + best_loss = val_epoch_loss + not_improved_count = 0 + best_state = True + else: + not_improved_count += 1 + best_state = False + # early stop + if self.args.early_stop: + if not_improved_count == self.args.early_stop_patience: + self.logger.info("Validation performance didn\'t improve for {} epochs. " + "Training stops.".format(self.args.early_stop_patience)) + break + # save the best state + if best_state == True: + self.logger.info('*********************************Current best model saved!') + best_model = copy.deepcopy(self.model.state_dict()) + + # if epoch%10==0:#test + # self.model.load_state_dict(best_model) + # #self.val_epoch(self.args.epochs, self.test_loader) + # self.test_simple(self.model, self.args, self.test_loader, self.scaler, self.logger, None, self.times) + + training_time = time.time() - start_time + self.logger.info("Total training time: {:.4f}min, best loss: {:.6f}".format((training_time / 60), best_loss)) + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass + + #save the best model to file + if not self.args.debug: + torch.save(best_model, self.best_path) + self.logger.info("Saving current best model to " + self.best_path) + + # if epoch==10:#test + # self.model.load_state_dict(best_model) + # #self.val_epoch(self.args.epochs, self.test_loader) + # self.test(self.model, self.args, self.test_loader, self.scaler, self.logger, None, self.times) + self.model.load_state_dict(best_model) + #self.val_epoch(self.args.epochs, self.test_loader) + self.test(self.model, self.args, self.test_loader, self.scaler, self.logger, None, self.times) + + def save_checkpoint(self): + state = { + 'state_dict': self.model.state_dict(), + 'optimizer': self.optimizer.state_dict(), + 'config': self.args + } + torch.save(state, self.best_path) + self.logger.info("Saving current best model to " + self.best_path) + + @staticmethod + def test(model, args, data_loader, scaler, logger, path, times): + if path != None: + check_point = torch.load(path) + state_dict = check_point['state_dict'] + args = check_point['config'] + model.load_state_dict(state_dict) + model.to(args.device) + model.eval() + y_pred = [] + y_true = [] + with torch.no_grad(): + for batch_idx, batch in enumerate(data_loader): + batch = tuple(b.to(args.device, dtype=torch.float) for b in batch) + *test_coeffs, target = batch + label = target[..., :args.output_dim] + output = model(times.to(args.device, dtype=torch.float), test_coeffs) + y_true.append(label) + y_pred.append(output) + y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) + if args.real_value: + y_pred = torch.cat(y_pred, dim=0) + else: + y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + np.save(args.log_dir+'/{}_true.npy'.format(args.dataset), y_true.cpu().numpy()) + np.save(args.log_dir+'/{}_pred.npy'.format(args.dataset), y_pred.cpu().numpy()) + for t in range(y_true.shape[1]): + mae, rmse, mape, _, _ = All_Metrics(y_pred[:, t, ...], y_true[:, t, ...], + args.mae_thresh, args.mape_thresh) + logger.info("Horizon {:02d}, MAE: {:.2f}, RMSE: {:.2f}, MAPE: {:.4f}%".format( + t + 1, mae, rmse, mape*100)) + mae, rmse, mape, _, _ = All_Metrics(y_pred, y_true, args.mae_thresh, args.mape_thresh) + logger.info("Average Horizon, MAE: {:.2f}, RMSE: {:.2f}, MAPE: {:.4f}%".format( + mae, rmse, mape*100)) + + @staticmethod + def test_simple(model, args, data_loader, scaler, logger, path, times): + if path != None: + check_point = torch.load(path) + state_dict = check_point['state_dict'] + args = check_point['config'] + model.load_state_dict(state_dict) + model.to(args.device) + model.eval() + y_pred = [] + y_true = [] + with torch.no_grad(): + for batch_idx, batch in enumerate(data_loader): + # for batch_idx, (data, target) in enumerate(data_loader): + batch = tuple(b.to(args.device, dtype=torch.float) for b in batch) + *test_coeffs, target = batch + # data = data[..., :args.input_dim] + label = target[..., :args.output_dim] + output = model(times.to(args.device, dtype=torch.float), test_coeffs) + y_true.append(label) + y_pred.append(output) + y_true = scaler.inverse_transform(torch.cat(y_true, dim=0)) + if args.real_value: + y_pred = torch.cat(y_pred, dim=0) + else: + y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) + + for t in range(y_true.shape[1]): + mae, rmse, mape, _, _ = All_Metrics(y_pred[:, t, ...], y_true[:, t, ...], + args.mae_thresh, args.mape_thresh) + mae, rmse, mape, _, _ = All_Metrics(y_pred, y_true, args.mae_thresh, args.mape_thresh) + logger.info("Average Horizon, MAE: {:.2f}, RMSE: {:.2f}, MAPE: {:.4f}%".format( + mae, rmse, mape*100)) + + @staticmethod + def _compute_sampling_threshold(global_step, k): + """ + Computes the sampling probability for scheduled sampling using inverse sigmoid. + :param global_step: + :param k: + :return: + """ + return k / (k + math.exp(global_step / k)) + +def _add_weight_regularisation(total_loss, regularise_parameters, scaling=0.03): + for parameter in regularise_parameters.parameters(): + if parameter.requires_grad: + total_loss = total_loss + scaling * parameter.norm() + return total_loss \ No newline at end of file diff --git a/model/STGNRDE/GRDE.py b/model/STGNRDE/GRDE.py new file mode 100755 index 0000000..4b64d8f --- /dev/null +++ b/model/STGNRDE/GRDE.py @@ -0,0 +1,120 @@ +import torch +import torch.nn.functional as F +import torch.nn as nn + +import model.STGNRDE.torchcde as torchcde +from model.STGNRDE.vector_fields import * + +class NeuralGCDE(nn.Module): + def __init__(self, args, func_f, func_g, input_channels, hidden_channels, output_channels, initial, device, atol, rtol, solver): + super(NeuralGCDE, self).__init__() + self.num_node = args['num_nodes'] + self.input_dim = input_channels + self.hidden_dim = hidden_channels + self.output_dim = output_channels + self.horizon = args['horizon'] + self.num_layers = args['num_layers'] + + # defaults for NRDE runtime options + self.model_type = args.get('model_type', 'rde') # {'rde', 'rde2'} + self.emb_opt = args.get('emb_opt', False) + self.interpolation = args.get('interpolation', 'cubic') # {'cubic','linear'} + self.adp_opt = args.get('adp_opt', False) + + self.default_graph = args['default_graph'] + self.node_embeddings = nn.Parameter(torch.randn(self.num_node, args['embed_dim']), requires_grad=True) + + self.func_f = func_f + self.func_g = func_g + self.solver = solver + self.atol = atol + self.rtol = rtol + + #predictor + self.end_conv = nn.Conv2d(1, args['horizon'] * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True) + + self.init_type = 'fc' + if self.init_type == 'fc': + self.initial_h = torch.nn.Linear(self.input_dim, self.hidden_dim) + self.initial_z = torch.nn.Linear(self.input_dim, self.hidden_dim) + elif self.init_type == 'conv': + self.start_conv_h = nn.Conv2d(in_channels=input_channels, + out_channels=hidden_channels, + kernel_size=(1,1)) + self.start_conv_z = nn.Conv2d(in_channels=input_channels, + out_channels=hidden_channels, + kernel_size=(1,1)) + + # optional projection for adaptive pooling path + if self.adp_opt: + self.proj = nn.Linear(self.hidden_dim, 1) + + def red_emb(self, coeffs): + # no-op reduction by default + return coeffs + + def forward(self, times, coeffs): + # source: B, T_1, N, D + # target: B, T_2, N, D + # times = torch.linspace(0, len(times)-1, coeffs.size(-2)).to(coeffs.device) + # times = torch.linspace(0, coeffs.size(-1), coeffs.size(-2)).to(coeffs.device) + # times = torch.linspace(0, coeffs.size(-2)-1, coeffs.size(-2)).to(coeffs.device) + # import pdb; pdb.set_trace() + if self.emb_opt == True: + coeffs = self.red_emb(coeffs) + # Adapt coeffs from tuple/list (a, b, two_c, three_d) to concatenated tensor if necessary + if isinstance(coeffs, (list, tuple)): + coeffs = torch.cat(coeffs, dim=-1) + if self.interpolation == 'cubic': + X = torchcde.CubicSpline(coeffs) + elif self.interpolation == 'linear': + X = torchcde.LinearInterpolation(coeffs) + X0 = X.evaluate(X.interval[0]) + + if self.init_type == 'fc': + h0 = self.initial_h(X0) + z0 = self.initial_z(X0) + # z0 = self.initial_z(h0) + elif self.init_type == 'conv': + h0 = self.start_conv_h(X0.transpose(1, 2).unsqueeze(-1)).transpose(1, 2).squeeze() + z0 = self.start_conv_z(X0.transpose(1, 2).unsqueeze(-1)).transpose(1, 2).squeeze() + + if self.model_type == 'rde': + z_T = torchcde.cdeint(X=X, + func=self.func_g, + z0=z0, + t=times, + adjoint=True, + method=self.solver + ) + elif self.model_type == 'rde2': + step_size = (X.grid_points[1:] - X.grid_points[:-1]).min() + # adjoint_params = tuple(self.func_f.parameters()) + tuple(self.func_g.parameters()) + (coeffs,) + + z_T = torchcde.cdeint_custom(X=X, + func_f=self.func_f, + func_g=self.func_g, + h0=h0, + z0=z0, + t=times, + adjoint=True, + method=self.solver, + ) + + if self.adp_opt == False: + z_T = z_T[:, :, -1:, :].transpose(1, 2) + else: + z_T = z_T.transpose(1, 2) + retain_score = self.proj(z_T) + retain_score = retain_score.squeeze() + retain_score = torch.sigmoid(retain_score.transpose(-1, -2)) + retain_score = retain_score.unsqueeze(-1) + z_T = torch.matmul(retain_score.transpose(-1, -2), z_T.permute(0, 2, 1, 3)).transpose(1, 2) + + # CNN based predictor + # output = self.end_conv(z_T.unsqueeze(-1).shape) #B, T*C, N, 1 + output = self.end_conv(z_T) # B, T*C, N, 1 + output = output.squeeze(-1).reshape(-1, self.horizon, self.output_dim, self.num_node) + output = output.permute(0, 1, 3, 2) # B, T, N, C + + return output \ No newline at end of file diff --git a/model/STGNRDE/Make_model.py b/model/STGNRDE/Make_model.py new file mode 100755 index 0000000..627426c --- /dev/null +++ b/model/STGNRDE/Make_model.py @@ -0,0 +1,36 @@ +from model.STGNRDE.GRDE import NeuralGCDE +from model.STGNRDE.vector_fields import FinalTanh_f, VectorField_g + + +def make_model(args): + vector_field_f = FinalTanh_f( + input_channels=args['input_dim'], + hidden_channels=args['hid_dim'], + hidden_hidden_channels=args['hid_hid_dim'], + num_hidden_layers=args['num_layers'] + ) + vector_field_g = VectorField_g( + input_channels=args['input_dim'], + hidden_channels=args['hid_dim'], + hidden_hidden_channels=args['hid_hid_dim'], + num_hidden_layers=args['num_layers'], + num_nodes=args['num_nodes'], + cheb_k=args['cheb_k'], + embed_dim=args['embed_dim'], + g_type=args['g_type'] + ) + model = NeuralGCDE( + args, + func_f=vector_field_f, + func_g=vector_field_g, + input_channels=args['input_dim'], + hidden_channels=args['hid_dim'], + output_channels=args['output_dim'], + initial=True, + device=args['device'], + atol=1e-9, + rtol=1e-7, + solver=args['solver'] + ) + return model + diff --git a/model/STGNRDE/PEMSD4_GCDE.conf b/model/STGNRDE/PEMSD4_GCDE.conf new file mode 100755 index 0000000..b4deda7 --- /dev/null +++ b/model/STGNRDE/PEMSD4_GCDE.conf @@ -0,0 +1,45 @@ +[data] +num_nodes = 307 +lag = 12 +horizon = 12 +val_ratio = 0.2 +test_ratio = 0.2 +tod = False +normalizer = std +column_wise = False +default_graph = True + +[model] +type = type1 +g_type = agc +input_dim = 2 +output_dim = 1 +embed_dim = 10 +hid_dim = 128 +hid_hid_dim = 128 +num_layers = 3 +cheb_order = 2 + +[train] +loss_func = mae +seed = 10 +batch_size = 64 +epochs = 100 +lr_init = 0.001 +weight_decay = 1e-3 +lr_decay = False +lr_decay_rate = 0.3 +lr_decay_step = 5,20,40,70 +early_stop = True +early_stop_patience = 15 +grad_norm = False +max_grad_norm = 5 +real_value = True + +[test] +mae_thresh = None +mape_thresh = 0. + +[log] +log_step = 20 +plot = False \ No newline at end of file diff --git a/model/STGNRDE/Run_cde.py b/model/STGNRDE/Run_cde.py new file mode 100755 index 0000000..1a20ac8 --- /dev/null +++ b/model/STGNRDE/Run_cde.py @@ -0,0 +1,205 @@ +import os +import sys +file_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +print(file_dir) +sys.path.append(file_dir) + +import torch +import numpy as np +import torch.nn as nn +import argparse +import configparser +import time + +from model.BasicTrainer_cde import Trainer +from lib.TrainInits import init_seed +from lib.dataloader import get_dataloader_cde +from lib.TrainInits import print_model_parameters +import os +from os.path import join +from Make_model import make_model +from torch.utils.tensorboard import SummaryWriter +#*************************************************************************# +Mode = 'train' +DEBUG = 'False' +DATASET = 'PEMSD4' #PEMSD4 or PEMSD8 +MODEL = 'GCDE' + +#get configuration +config_file = './{}_{}.conf'.format(DATASET, MODEL) +#print('Read configuration file: %s' % (config_file)) +config = configparser.ConfigParser() +config.read(config_file) + +from lib.metrics import MAE_torch +def masked_mae_loss(scaler, mask_value): + def loss(preds, labels): + if scaler: + preds = scaler.inverse_transform(preds) + labels = scaler.inverse_transform(labels) + mae = MAE_torch(pred=preds, true=labels, mask_value=mask_value) + return mae + return loss + +#parser +args = argparse.ArgumentParser(description='arguments') +args.add_argument('--dataset', default=DATASET, type=str) +args.add_argument('--mode', default=Mode, type=str) +args.add_argument('--device', default=0, type=int, help='indices of GPUs') +args.add_argument('--debug', default=DEBUG, type=eval) +args.add_argument('--model', default=MODEL, type=str) +args.add_argument('--cuda', default=True, type=bool) +args.add_argument('--comment', default='', type=str) + + +#data +args.add_argument('--val_ratio', default=config['data']['val_ratio'], type=float) +args.add_argument('--test_ratio', default=config['data']['test_ratio'], type=float) +args.add_argument('--lag', default=config['data']['lag'], type=int) +args.add_argument('--horizon', default=config['data']['horizon'], type=int) +args.add_argument('--num_nodes', default=config['data']['num_nodes'], type=int) +args.add_argument('--tod', default=config['data']['tod'], type=eval) +args.add_argument('--normalizer', default=config['data']['normalizer'], type=str) +args.add_argument('--column_wise', default=config['data']['column_wise'], type=eval) +args.add_argument('--default_graph', default=config['data']['default_graph'], type=eval) +#model +args.add_argument('--model_type', default=config['model']['type'], type=str) +args.add_argument('--g_type', default=config['model']['g_type'], type=str) +args.add_argument('--input_dim', default=config['model']['input_dim'], type=int) +args.add_argument('--output_dim', default=config['model']['output_dim'], type=int) +args.add_argument('--embed_dim', default=config['model']['embed_dim'], type=int) +args.add_argument('--hid_dim', default=config['model']['hid_dim'], type=int) +args.add_argument('--hid_hid_dim', default=config['model']['hid_hid_dim'], type=int) +args.add_argument('--num_layers', default=config['model']['num_layers'], type=int) +args.add_argument('--cheb_k', default=config['model']['cheb_order'], type=int) +args.add_argument('--solver', default='rk4', type=str) + +#train +args.add_argument('--loss_func', default=config['train']['loss_func'], type=str) +args.add_argument('--seed', default=config['train']['seed'], type=int) +args.add_argument('--batch_size', default=config['train']['batch_size'], type=int) +args.add_argument('--epochs', default=config['train']['epochs'], type=int) +args.add_argument('--lr_init', default=config['train']['lr_init'], type=float) +args.add_argument('--weight_decay', default=config['train']['weight_decay'], type=eval) +args.add_argument('--lr_decay', default=config['train']['lr_decay'], type=eval) +args.add_argument('--lr_decay_rate', default=config['train']['lr_decay_rate'], type=float) +args.add_argument('--lr_decay_step', default=config['train']['lr_decay_step'], type=str) +args.add_argument('--early_stop', default=config['train']['early_stop'], type=eval) +args.add_argument('--early_stop_patience', default=config['train']['early_stop_patience'], type=int) +args.add_argument('--grad_norm', default=config['train']['grad_norm'], type=eval) +args.add_argument('--max_grad_norm', default=config['train']['max_grad_norm'], type=int) +args.add_argument('--teacher_forcing', default=False, type=bool) +#args.add_argument('--tf_decay_steps', default=2000, type=int, help='teacher forcing decay steps') +args.add_argument('--real_value', default=config['train']['real_value'], type=eval, help = 'use real value for loss calculation') + +args.add_argument('--missing_test', default=False, type=bool) +args.add_argument('--missing_rate', default=0.1, type=float) + +#test +args.add_argument('--mae_thresh', default=config['test']['mae_thresh'], type=eval) +args.add_argument('--mape_thresh', default=config['test']['mape_thresh'], type=float) +args.add_argument('--model_path', default='', type=str) +#log +args.add_argument('--log_dir', default='../runs', type=str) +args.add_argument('--log_step', default=config['log']['log_step'], type=int) +args.add_argument('--plot', default=config['log']['plot'], type=eval) +args.add_argument('--tensorboard',action='store_true',help='tensorboard') + +args = args.parse_args() +init_seed(args.seed) + +GPU_NUM = args.device +device = torch.device(f'cuda:{GPU_NUM}' if torch.cuda.is_available() else 'cpu') +torch.cuda.set_device(device) # change allocation of current GPU + +print(args) + +#config log path +save_name = time.strftime("%m-%d-%Hh%Mm")+args.comment+"_"+ args.dataset+"_"+ args.model+"_"+ args.model_type+"_"+"embed{"+str(args.embed_dim)+"}"+"hid{"+str(args.hid_dim)+"}"+"hidhid{"+str(args.hid_hid_dim)+"}"+"lyrs{"+str(args.num_layers)+"}"+"lr{"+str(args.lr_init)+"}"+"wd{"+str(args.weight_decay)+"}" +path = '../runs' + +log_dir = join(path, args.dataset, save_name) +args.log_dir = log_dir +if (os.path.exists(args.log_dir)): + print('has model save path') +else: + os.makedirs(args.log_dir) + +if args.tensorboard: + w : SummaryWriter = SummaryWriter(args.log_dir) +else: + w = None + +#init model +if args.model_type=='type1': + model, vector_field_f, vector_field_g = make_model(args) +elif args.model_type=='type1_temporal': + model, vector_field_f = make_model(args) +elif args.model_type=='type1_spatial': + model, vector_field_g = make_model(args) +else: + raise ValueError("Check args.model_type") + +model = model.to(args.device) + +if args.model_type=='type1_temporal': + vector_field_f = vector_field_f.to(args.device) + vector_field_g = None +elif args.model_type=='type1_spatial': + vector_field_f = None + vector_field_g = vector_field_g.to(args.device) +else: + vector_field_f = vector_field_f.to(args.device) + vector_field_g = vector_field_g.to(args.device) + +print(model) + +for p in model.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + else: + nn.init.uniform_(p) +print_model_parameters(model, only_num=False) + +#load dataset +train_loader, val_loader, test_loader, scaler, times = get_dataloader_cde(args, + normalizer=args.normalizer, + tod=args.tod, dow=False, + weather=False, single=False) + +#init loss function, optimizer +if args.loss_func == 'mask_mae': + loss = masked_mae_loss(scaler, mask_value=0.0) +elif args.loss_func == 'mae': + loss = torch.nn.L1Loss().to(args.device) +elif args.loss_func == 'mse': + loss = torch.nn.MSELoss().to(args.device) +elif args.loss_func == 'huber_loss': + loss = torch.nn.HuberLoss(delta=1.0).to(args.device) +else: + raise ValueError + +optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr_init, + weight_decay=args.weight_decay) + +#learning rate decay +lr_scheduler = None +if args.lr_decay: + print('Applying learning rate decay.') + lr_decay_steps = [int(i) for i in list(args.lr_decay_step.split(','))] + lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer=optimizer, + milestones=lr_decay_steps, + gamma=args.lr_decay_rate) + +#start training +trainer = Trainer(model, vector_field_f, vector_field_g, loss, optimizer, train_loader, val_loader, test_loader, scaler, + args, lr_scheduler, args.device, times, + w) +if args.mode == 'train': + trainer.train() +elif args.mode == 'test': + model.load_state_dict(torch.load('./pre-trained/{}.pth'.format(args.dataset))) + print("Load saved model") + trainer.test(model, trainer.args, test_loader, scaler, trainer.logger, times) +else: + raise ValueError diff --git a/model/STGNRDE/torchcde/__init__.py b/model/STGNRDE/torchcde/__init__.py new file mode 100644 index 0000000..4fd7b84 --- /dev/null +++ b/model/STGNRDE/torchcde/__init__.py @@ -0,0 +1,9 @@ +from .interpolation_base import InterpolationBase +from .interpolation_cubic import natural_cubic_spline_coeffs, natural_cubic_coeffs, CubicSpline +from .interpolation_linear import linear_interpolation_coeffs, LinearInterpolation +from .interpolation_hermite_cubic_bdiff import hermite_cubic_coefficients_with_backward_differences +from .log_ode import logsignature_windows, logsig_windows +from .misc import TupleControl +from .solver import cdeint, cdeint_custom + +__version__ = "0.2.3" diff --git a/model/STGNRDE/torchcde/interpolation_base.py b/model/STGNRDE/torchcde/interpolation_base.py new file mode 100644 index 0000000..afc6ea4 --- /dev/null +++ b/model/STGNRDE/torchcde/interpolation_base.py @@ -0,0 +1,22 @@ +import abc +import torch + + +class InterpolationBase(torch.nn.Module, metaclass=abc.ABCMeta): + @property + @abc.abstractmethod + def grid_points(self): + raise NotImplementedError + + @property + @abc.abstractmethod + def interval(self): + raise NotImplementedError + + @abc.abstractmethod + def evaluate(self, t): + raise NotImplementedError + + @abc.abstractmethod + def derivative(self, t): + raise NotImplementedError diff --git a/model/STGNRDE/torchcde/interpolation_cubic.py b/model/STGNRDE/torchcde/interpolation_cubic.py new file mode 100644 index 0000000..976126e --- /dev/null +++ b/model/STGNRDE/torchcde/interpolation_cubic.py @@ -0,0 +1,346 @@ +import torch +from torchcde import interpolation_base + +from . import misc + + +def _natural_cubic_spline_coeffs_without_missing_values(t, x): + # x should be a tensor of shape (..., length) + # Will return the b, two_c, three_d coefficients of the derivative of the cubic spline interpolating the path. + + length = x.size(-1) + + if length < 2: + # In practice this should always already be caught in __init__. + raise ValueError("Must have a time dimension of size at least 2.") + elif length == 2: + a = x[..., :1] + b = (x[..., 1:] - x[..., :1]) / (t[..., 1:] - t[..., :1]) + two_c = torch.zeros(*x.shape[:-1], 1, dtype=x.dtype, device=x.device) + three_d = torch.zeros(*x.shape[:-1], 1, dtype=x.dtype, device=x.device) + else: + # Set up some intermediate values + time_diffs = t[1:] - t[:-1] + time_diffs_reciprocal = time_diffs.reciprocal() + time_diffs_reciprocal_squared = time_diffs_reciprocal ** 2 + three_path_diffs = 3 * (x[..., 1:] - x[..., :-1]) + six_path_diffs = 2 * three_path_diffs + path_diffs_scaled = three_path_diffs * time_diffs_reciprocal_squared + + # Solve a tridiagonal linear system to find the derivatives at the knots + system_diagonal = torch.empty(length, dtype=x.dtype, device=x.device) + system_diagonal[:-1] = time_diffs_reciprocal + system_diagonal[-1] = 0 + system_diagonal[1:] += time_diffs_reciprocal + system_diagonal *= 2 + system_rhs = torch.empty_like(x) + system_rhs[..., :-1] = path_diffs_scaled + system_rhs[..., -1] = 0 + system_rhs[..., 1:] += path_diffs_scaled + knot_derivatives = misc.tridiagonal_solve(system_rhs, time_diffs_reciprocal, system_diagonal, + time_diffs_reciprocal) + + # Do some algebra to find the coefficients of the spline + a = x[..., :-1] + b = knot_derivatives[..., :-1] + two_c = (six_path_diffs * time_diffs_reciprocal + - 4 * knot_derivatives[..., :-1] + - 2 * knot_derivatives[..., 1:]) * time_diffs_reciprocal + three_d = (-six_path_diffs * time_diffs_reciprocal + + 3 * (knot_derivatives[..., :-1] + + knot_derivatives[..., 1:])) * time_diffs_reciprocal_squared + + return a, b, two_c, three_d + + +def _natural_cubic_spline_coeffs_with_missing_values(t, x, _version): + if x.ndimension() == 1: + # We have to break everything down to individual scalar paths because of the possibility of missing values + # being different in different channels + return _natural_cubic_spline_coeffs_with_missing_values_scalar(t, x, _version) + else: + a_pieces = [] + b_pieces = [] + two_c_pieces = [] + three_d_pieces = [] + for p in x.unbind(dim=0): # TODO: parallelise over this + a, b, two_c, three_d = _natural_cubic_spline_coeffs_with_missing_values(t, p, _version) + a_pieces.append(a) + b_pieces.append(b) + two_c_pieces.append(two_c) + three_d_pieces.append(three_d) + return (misc.cheap_stack(a_pieces, dim=0), + misc.cheap_stack(b_pieces, dim=0), + misc.cheap_stack(two_c_pieces, dim=0), + misc.cheap_stack(three_d_pieces, dim=0)) + + +def _natural_cubic_spline_coeffs_with_missing_values_scalar(t, x, _version): + # t and x both have shape (length,) + + nan = torch.isnan(x) + not_nan = ~nan + path_no_nan = x.masked_select(not_nan) + + if path_no_nan.size(0) == 0: + # Every entry is a NaN, so we take a constant path with derivative zero, so return zero coefficients. + # Note that we may assume that X.size(0) >= 2 by the checks in __init__ so "X.size(0) - 1" is a valid + # thing to do. + return (torch.zeros(x.size(0) - 1, dtype=x.dtype, device=x.device), + torch.zeros(x.size(0) - 1, dtype=x.dtype, device=x.device), + torch.zeros(x.size(0) - 1, dtype=x.dtype, device=x.device), + torch.zeros(x.size(0) - 1, dtype=x.dtype, device=x.device)) + # else we have at least one non-NaN entry, in which case we're going to impute at least one more entry (as + # the path is of length at least 2 so the start and the end aren't the same), so we will then have at least two + # non-Nan entries. In particular we can call _compute_coeffs safely later. + + # How to deal with missing values at the start or end of the time series? We're creating some splines, so one + # option is just to extend the first piece backwards, and the final piece forwards. But polynomials tend to + # behave badly when extended beyond the interval they were constructed on, so the results can easily end up + # being awful. + if _version == 0: + # Instead we impute an observation at the very start equal to the first actual observation made, and impute an + # observation at the very end equal to the last actual observation made, and then proceed with splines as + # normal. + need_new_not_nan = False + if torch.isnan(x[0]): + if not need_new_not_nan: + x = x.clone() + need_new_not_nan = True + x[0] = path_no_nan[0] + if torch.isnan(x[-1]): + if not need_new_not_nan: + x = x.clone() + need_new_not_nan = True + x[-1] = path_no_nan[-1] + if need_new_not_nan: + not_nan = ~torch.isnan(x) + path_no_nan = x.masked_select(not_nan) + else: + # Instead we fill forward and backward from the first/last observation made. This is better than the previous + # approach as the splines instead rapidly stabilise to the first/last value. + cumsum_mask = not_nan.cumsum(dim=0) + cumsum_mask[nan] = -1 + last_non_nan_index = cumsum_mask.argmax(dim=0) + cumsum_mask[nan] = 1 + last_non_nan_index + first_non_nan_index = cumsum_mask.argmin(dim=0) + x = x.clone() + x[:first_non_nan_index] = x[first_non_nan_index] + x[last_non_nan_index + 1:] = x[last_non_nan_index] + not_nan = ~torch.isnan(x) + path_no_nan = x.masked_select(not_nan) + times_no_nan = t.masked_select(not_nan) + + # Find the coefficients on the pieces we do understand + # These all have shape (len - 1,) + (a_pieces_no_nan, + b_pieces_no_nan, + two_c_pieces_no_nan, + three_d_pieces_no_nan) = _natural_cubic_spline_coeffs_without_missing_values(times_no_nan, path_no_nan) + + # Now we're going to normalise them to give coefficients on every interval + a_pieces = [] + b_pieces = [] + two_c_pieces = [] + three_d_pieces = [] + + iter_times_no_nan = iter(times_no_nan) + iter_coeffs_no_nan = iter(zip(a_pieces_no_nan, b_pieces_no_nan, two_c_pieces_no_nan, three_d_pieces_no_nan)) + next_time_no_nan = next(iter_times_no_nan) + for time in t[:-1]: + # will always trigger on the first iteration because of how we've imputed missing values at the start and + # end of the time series. + if time >= next_time_no_nan: + prev_time_no_nan = next_time_no_nan + next_time_no_nan = next(iter_times_no_nan) + next_a_no_nan, next_b_no_nan, next_two_c_no_nan, next_three_d_no_nan = next(iter_coeffs_no_nan) + offset = prev_time_no_nan - time + a_inner = (0.5 * next_two_c_no_nan - next_three_d_no_nan * offset / 3) * offset + a_pieces.append(next_a_no_nan + (a_inner - next_b_no_nan) * offset) + b_pieces.append(next_b_no_nan + (next_three_d_no_nan * offset - next_two_c_no_nan) * offset) + two_c_pieces.append(next_two_c_no_nan - 2 * next_three_d_no_nan * offset) + three_d_pieces.append(next_three_d_no_nan) + + return (misc.cheap_stack(a_pieces, dim=0), + misc.cheap_stack(b_pieces, dim=0), + misc.cheap_stack(two_c_pieces, dim=0), + misc.cheap_stack(three_d_pieces, dim=0)) + + +# The mathematics of this are adapted from http://mathworld.wolfram.com/CubicSpline.html, although they only treat the +# case of each piece being parameterised by [0, 1]. (We instead take the length of each piece to be the difference in +# time stamps.) +def _natural_cubic_spline_coeffs(x, t, _version): + t = misc.validate_input_path(x, t) + + if torch.isnan(x).any(): + # Transpose because channels are a batch dimension for the purpose of finding interpolating polynomials. + # b, two_c, three_d have shape (..., channels, length - 1) + a, b, two_c, three_d = _natural_cubic_spline_coeffs_with_missing_values(t, x.transpose(-1, -2), _version) + else: + # Can do things more quickly in this case. + a, b, two_c, three_d = _natural_cubic_spline_coeffs_without_missing_values(t, x.transpose(-1, -2)) + + # These all have shape (..., length - 1, channels) + a = a.transpose(-1, -2) + b = b.transpose(-1, -2) + two_c = two_c.transpose(-1, -2) + three_d = three_d.transpose(-1, -2) + coeffs = torch.cat([a, b, two_c, three_d], dim=-1) # for simplicity put them all together + return coeffs + + +def natural_cubic_spline_coeffs(x, t=None): + """Calculates the coefficients of the natural cubic spline approximation to the batch of controls given. + + ******************** + DEPRECATED: this now exists for backward compatibility. For new projects please use `natural_cubic_coeffs` instead, + which handles missing data at the start/end of a time series better. + ******************** + + Arguments: + x: tensor of values, of shape (..., length, input_channels), where ... is some number of batch dimensions. This + is interpreted as a (batch of) paths taking values in an input_channels-dimensional real vector space, with + length-many observations. Missing values are supported, and should be represented as NaNs. + t: Optional one dimensional tensor of times. Must be monotonically increasing. If not passed will default to + tensor([0., 1., ..., length - 1]). If you are using neural CDEs then you **do not need to use this + argument**. See the Further Documentation in README.md. + + Warning: + If there are missing values then calling this function can be pretty slow. Make sure to cache the result, and + don't reinstantiate it on every forward pass, if at all possible. + + Returns: + A tensor, which should in turn be passed to `torchcde.CubicSpline`. + + Why do we do it like this? Because typically you want to use PyTorch tensors at various interfaces, for example + when loading a batch from a DataLoader. If we wrapped all of this up into just the + `torchcde.CubicSpline` class then that sort of thing wouldn't be possible. + + As such the suggested use is to: + (a) Load your data. + (b) Preprocess it with this function. + (c) Save the result. + (d) Treat the result as your dataset as far as PyTorch's `torch.utils.data.Dataset` and + `torch.utils.data.DataLoader` classes are concerned. + (e) Call CubicSpline as the first part of your model. + + See also the accompanying example.py. + """ + return _natural_cubic_spline_coeffs(x, t, _version=0) + + +def natural_cubic_coeffs(x, t=None): + """Calculates the coefficients of the natural cubic spline approximation to the batch of controls given. + + Arguments: + x: tensor of values, of shape (..., length, input_channels), where ... is some number of batch dimensions. This + is interpreted as a (batch of) paths taking values in an input_channels-dimensional real vector space, with + length-many observations. Missing values are supported, and should be represented as NaNs. + t: Optional one dimensional tensor of times. Must be monotonically increasing. If not passed will default to + tensor([0., 1., ..., length - 1]). If you are using neural CDEs then you **do not need to use this + argument**. See the Further Documentation in README.md. + + Warning: + If there are missing values then calling this function can be pretty slow. Make sure to cache the result, and + don't reinstantiate it on every forward pass, if at all possible. + + Returns: + A tensor, which should in turn be passed to `torchcde.CubicSpline`. + + Why do we do it like this? Because typically you want to use PyTorch tensors at various interfaces, for example + when loading a batch from a DataLoader. If we wrapped all of this up into just the + `torchcde.CubicSpline` class then that sort of thing wouldn't be possible. + + As such the suggested use is to: + (a) Load your data. + (b) Preprocess it with this function. + (c) Save the result. + (d) Treat the result as your dataset as far as PyTorch's `torch.utils.data.Dataset` and + `torch.utils.data.DataLoader` classes are concerned. + (e) Call CubicSpline as the first part of your model. + + See also the accompanying example.py. + """ + return _natural_cubic_spline_coeffs(x, t, _version=1) + + +class CubicSpline(interpolation_base.InterpolationBase): + """Calculates the cubic spline approximation to the batch of controls given. Also calculates its derivative. + + Example: + # (2, 1) are batch dimensions. 7 is the time dimension (of the same length as t). 3 is the channel dimension. + x = torch.rand(2, 1, 7, 3) + coeffs = natural_cubic_coeffs(x) + # ...at this point you can save coeffs, put it through PyTorch's Datasets and DataLoaders, etc... + spline = CubicSpline(coeffs) + point = torch.tensor(0.4) + # will be a tensor of shape (2, 1, 3), corresponding to batch and channel dimensions + out = spline.derivative(point) + """ + + def __init__(self, coeffs, t=None, **kwargs): + """ + Arguments: + coeffs: As returned by `torchcde.natural_cubic_coeffs`. + t: As passed to linear_interpolation_coeffs. (If it was passed. If you are using neural CDEs then you **do + not need to use this argument**. See the Further Documentation in README.md.) + """ + super(CubicSpline, self).__init__(**kwargs) + + if t is None: + t = torch.linspace(0, coeffs.size(-2), coeffs.size(-2) + 1, dtype=coeffs.dtype, device=coeffs.device) + + channels = coeffs.size(-1) // 4 + if channels * 4 != coeffs.size(-1): # check that it's a multiple of 4 + raise ValueError("Passed invalid coeffs.") + a, b, two_c, three_d = (coeffs[..., :channels], coeffs[..., channels:2 * channels], + coeffs[..., 2 * channels:3 * channels], coeffs[..., 3 * channels:]) + + self.register_buffer('_t', t) + self.register_buffer('_a', a) + self.register_buffer('_b', b) + # as we're typically computing derivatives, we store the multiples of these coefficients that are more useful + self.register_buffer('_two_c', two_c) + self.register_buffer('_three_d', three_d) + + @property + def grid_points(self): + return self._t + + @property + def interval(self): + return torch.stack([self._t[0], self._t[-1]]) + + def _interpret_t(self, t): + t = torch.as_tensor(t, dtype=self._b.dtype, device=self._b.device) + maxlen = self._b.size(-2) - 1 + # clamp because t may go outside of [t[0], t[-1]]; this is fine + index = torch.bucketize(t.detach(), self._t.detach()).sub(1).clamp(0, maxlen) + # will never access the last element of self._t; this is correct behaviour + fractional_part = t - self._t[index] + return fractional_part, index + + def evaluate(self, t): + fractional_part, index = self._interpret_t(t) + fractional_part = fractional_part.unsqueeze(-1) + inner = 0.5 * self._two_c[..., index, :] + self._three_d[..., index, :] * fractional_part / 3 + inner = self._b[..., index, :] + inner * fractional_part + return self._a[..., index, :] + inner * fractional_part + + def derivative(self, t): + fractional_part, index = self._interpret_t(t) + fractional_part = fractional_part.unsqueeze(-1) + inner = self._two_c[..., index, :] + self._three_d[..., index, :] * fractional_part + deriv = self._b[..., index, :] + inner * fractional_part + return deriv + + +class NaturalCubicSpline(CubicSpline): + """Calculates the coefficients of the natural cubic spline approximation to the batch of controls given. + + ******************** + DEPRECATED: this now exists for backward compatibility. For new projects please use `CubicSpline` instead. This + class is general for any cubic coeffs (currently natural cubic or Hermite with backwards differences). + ******************** + """ diff --git a/model/STGNRDE/torchcde/interpolation_hermite_cubic_bdiff.py b/model/STGNRDE/torchcde/interpolation_hermite_cubic_bdiff.py new file mode 100644 index 0000000..a861081 --- /dev/null +++ b/model/STGNRDE/torchcde/interpolation_hermite_cubic_bdiff.py @@ -0,0 +1,44 @@ +import torch +from torchcde.interpolation_linear import linear_interpolation_coeffs + + +def _setup_hermite_cubic_coeffs_w_backward_differences(times, coeffs, derivs, device=None): + """Compute backward hermite from linear coeffs.""" + x_prev = coeffs[..., :-1, :] + x_next = coeffs[..., 1:, :] + # Let x_0 - x_{-1} = x_1 - x_0 + derivs_prev = torch.cat((derivs[..., [0], :], derivs[..., :-1, :]), axis=-2) + derivs_next = derivs + x_diff = x_next - x_prev + t_diff = (times[1:] - times[:-1]).unsqueeze(-1) + # Coeffs + a = x_prev + b = derivs_prev + two_c = 2 * (3 * (x_diff / t_diff - b) - derivs_next + derivs_prev) / t_diff + three_d = (1 / t_diff ** 2) * (derivs_next - b) - (two_c) / t_diff + coeffs = torch.cat([a, b, two_c, three_d], dim=-1).to(device) + return coeffs + + +def hermite_cubic_coefficients_with_backward_differences(x, t=None): + """Computes the coefficients for hermite cubic splines with backward differences. + + Arguments: + As `torchcde.linear_interpolation_coeffs`. + + Returns: + A tensor, which should in turn be passed to `torchcde.CubicSpline`. + """ + # Linear coeffs + coeffs = linear_interpolation_coeffs(x, t=t, rectilinear=None) + + if t is None: + t = torch.linspace(0, coeffs.size(-2) - 1, coeffs.size(-2), dtype=coeffs.dtype, device=coeffs.device) + + # Linear derivs + derivs = (coeffs[..., 1:, :] - coeffs[..., :-1, :]) / (t[1:] - t[:-1]).unsqueeze(-1) + + # Use the above to compute hermite coeffs + hermite_coeffs = _setup_hermite_cubic_coeffs_w_backward_differences(t, coeffs, derivs, device=coeffs.device) + + return hermite_coeffs diff --git a/model/STGNRDE/torchcde/interpolation_linear.py b/model/STGNRDE/torchcde/interpolation_linear.py new file mode 100644 index 0000000..f0c486a --- /dev/null +++ b/model/STGNRDE/torchcde/interpolation_linear.py @@ -0,0 +1,225 @@ +import math +import torch +import warnings + +from . import interpolation_base +from . import misc + + +_two_pi = 2 * math.pi +_inv_two_pi = 1 / _two_pi + + +def _linear_interpolation_coeffs_with_missing_values_scalar(t, x): + # t and X both have shape (length,) + + not_nan = ~torch.isnan(x) + path_no_nan = x.masked_select(not_nan) + + if path_no_nan.size(0) == 0: + # Every entry is a NaN, so we take a constant path with derivative zero, so return zero coefficients. + return torch.zeros(x.size(0), dtype=x.dtype, device=x.device) + + if path_no_nan.size(0) == x.size(0): + # Every entry is not-NaN, so just return. + return x + + x = x.clone() + # How to deal with missing values at the start or end of the time series? We impute an observation at the very start + # equal to the first actual observation made, and impute an observation at the very end equal to the last actual + # observation made, and then proceed as normal. + if torch.isnan(x[0]): + x[0] = path_no_nan[0] + if torch.isnan(x[-1]): + x[-1] = path_no_nan[-1] + + nan_indices = torch.arange(x.size(0), device=x.device).masked_select(torch.isnan(x)) + + if nan_indices.size(0) == 0: + # We only had missing values at the start or end + return x + + prev_nan_index = nan_indices[0] + prev_not_nan_index = prev_nan_index - 1 + prev_not_nan_indices = [prev_not_nan_index] + for nan_index in nan_indices[1:]: + if prev_nan_index != nan_index - 1: + prev_not_nan_index = nan_index - 1 + prev_nan_index = nan_index + prev_not_nan_indices.append(prev_not_nan_index) + + next_nan_index = nan_indices[-1] + next_not_nan_index = next_nan_index + 1 + next_not_nan_indices = [next_not_nan_index] + for nan_index in reversed(nan_indices[:-1]): + if next_nan_index != nan_index + 1: + next_not_nan_index = nan_index + 1 + next_nan_index = nan_index + next_not_nan_indices.append(next_not_nan_index) + next_not_nan_indices = reversed(next_not_nan_indices) + for prev_not_nan_index, nan_index, next_not_nan_index in zip(prev_not_nan_indices, + nan_indices, + next_not_nan_indices): + prev_stream = x[prev_not_nan_index] + next_stream = x[next_not_nan_index] + prev_time = t[prev_not_nan_index] + next_time = t[next_not_nan_index] + time = t[nan_index] + ratio = (time - prev_time) / (next_time - prev_time) + x[nan_index] = prev_stream + ratio * (next_stream - prev_stream) + + return x + + +def _linear_interpolation_coeffs_with_missing_values(t, x): + if x.ndimension() == 1: + # We have to break everything down to individual scalar paths because of the possibility of missing values + # being different in different channels + return _linear_interpolation_coeffs_with_missing_values_scalar(t, x) + else: + out_pieces = [] + for p in x.unbind(dim=0): # TODO: parallelise over this + out = _linear_interpolation_coeffs_with_missing_values(t, p) + out_pieces.append(out) + return misc.cheap_stack(out_pieces, dim=0) + + +def _prepare_rectilinear_interpolation(data, time_index): + """Prepares data for rectilinear interpolation. + + This function performs the relevant filling and lagging of the data needed to convert raw data into a format such + standard linear interpolation will give the rectilinear interpolation. + + Arguments: + x: tensor of values with first channel index being time, of shape (..., length, input_channels), where ... is + some number of batch dimensions. + time_index: integer giving the index of the time channel. + + Example: + Suppose we have data: + data = [(t1, x1), (t2, NaN), (t3, x3), ...] + that we wish to interpolate using a rectilinear scheme. The key point is that this is equivalent to a linear + interpolation on + data_rect = [(t1, x1), (t2, x1), (t2, x1), (t3, x1), (t3, x3) ...] + This function simply performs the conversion from `data` to `data_rect` so that we can apply the inbuilt + torchcde linear interpolation scheme to achieve rectilinear interpolation. + + Returns: + A tensor, now of shape (..., 2 * length - 1, input_channels] that can be fed to linear interpolation coeffs to + give rectilinear coeffs. + """ + # Check time_index is of the correct format + n_channels = data.size(-1) + assert isinstance(time_index, int), "Index of the time channel must be an integer in [0, {}]".format(n_channels - 1) + assert 0 <= time_index < n_channels, "Time index must be in [0, {}], was given {}." \ + "".format(n_channels - 1, time_index) + + times = data[..., time_index] + assert not torch.isnan(times).any(), "There exist nan values in the time column which is not allowed. If the " \ + "times are padded with nans after final time, a simple solution is to " \ + "forward fill the final time." + + # Forward fill and perform lag interleaving for rectilinear + data_filled = misc.forward_fill(data) + data_repeat = data_filled.repeat_interleave(2, dim=-2) + data_repeat[..., :-1, time_index] = data_repeat[..., 1:, time_index] + data_rect = data_repeat[..., :-1, :] + + return data_rect + + +def linear_interpolation_coeffs(x, t=None, rectilinear=None): + """Calculates the knots of the linear interpolation of the batch of controls given. + + Arguments: + x: tensor of values, of shape (..., length, input_channels), where ... is some number of batch dimensions. This + is interpreted as a (batch of) paths taking values in an input_channels-dimensional real vector space, with + length-many observations. Missing values are supported, and should be represented as NaNs. + t: Optional one dimensional tensor of times. Must be monotonically increasing. If not passed will default to + tensor([0., 1., ..., length - 1]). If you are using neural CDEs then you **do not need to use this + argument**. See the Further Documentation in README.md. + rectilinear: Optional integer. Used for performing rectilinear interpolation. This means that interpolation + between each two adjoint points is done by first interpolating in the time direction, and then interpolating + in the feature direction. (This is useful for causal missing data, see the Further Documentation in + README.md.) Defaults to None, i.e. not performing rectilinear interpolation. For rectilinear interpolation + time *must* be a channel in x and the `rectilinear` parameter must be an integer specifying the channel + index location of the time index in x. + + Warning: + If there are missing values then calling this function can be pretty slow. Make sure to cache the result, and + don't call it on every forward pass, if at all possible. + + Returns: + A tensor, which should in turn be passed to `torchcde.LinearInterpolation`. + + See the docstring for `torchcde.natural_cubic_coeffs` for more information on why we do it this way. + """ + if rectilinear is not None: + if torch.isnan(x[..., 0, :]).any(): + warnings.warn("The data `x` begins with missing values in some channels. The path will be constructed by " + "backward-filling the first observed value, which is not causal. Raising a warning as the " + "`rectilinear` argument has also been passed, which is nearly always only used when " + "causality is desired. If you need causality then fill in the missing value at the start of " + "each channel with whatever you'd like it to be. (The mean over that channel is a common " + "choice.)") + x = _prepare_rectilinear_interpolation(x, rectilinear) + + t = misc.validate_input_path(x, t) + + if torch.isnan(x).any(): + x = _linear_interpolation_coeffs_with_missing_values(t, x.transpose(-1, -2)).transpose(-1, -2) + return x + + +class LinearInterpolation(interpolation_base.InterpolationBase): + """Calculates the linear interpolation to the batch of controls given. Also calculates its derivative.""" + + def __init__(self, coeffs, t=None, **kwargs): + """ + Arguments: + coeffs: As returned by linear_interpolation_coeffs. + t: As passed to linear_interpolation_coeffs. (If it was passed. If you are using neural CDEs then you **do + not need to use this argument**. See the Further Documentation in README.md.) + """ + super(LinearInterpolation, self).__init__(**kwargs) + + if t is None: + t = torch.linspace(0, coeffs.size(-2) - 1, coeffs.size(-2), dtype=coeffs.dtype, device=coeffs.device) + + derivs = (coeffs[..., 1:, :] - coeffs[..., :-1, :]) / (t[1:] - t[:-1]).unsqueeze(-1) + + self.register_buffer('_t', t) + self.register_buffer('_coeffs', coeffs) + self.register_buffer('_derivs', derivs) + + @property + def grid_points(self): + return self._t + + @property + def interval(self): + return torch.stack([self._t[0], self._t[-1]]) + + def _interpret_t(self, t): + t = torch.as_tensor(t, dtype=self._derivs.dtype, device=self._derivs.device) + maxlen = self._derivs.size(-2) - 1 + # clamp because t may go outside of [t[0], t[-1]]; this is fine + index = torch.bucketize(t.detach(), self._t.detach()).sub(1).clamp(0, maxlen) + # will never access the last element of self._t; this is correct behaviour + fractional_part = t - self._t[index] + return fractional_part, index + + def evaluate(self, t): + fractional_part, index = self._interpret_t(t) + fractional_part = fractional_part.unsqueeze(-1) + prev_coeff = self._coeffs[..., index, :] + next_coeff = self._coeffs[..., index + 1, :] + prev_t = self._t[index] + next_t = self._t[index + 1] + diff_t = next_t - prev_t + return prev_coeff + fractional_part * (next_coeff - prev_coeff) / diff_t.unsqueeze(-1) + + def derivative(self, t): + fractional_part, index = self._interpret_t(t) + deriv = self._derivs[..., index, :] + return deriv diff --git a/model/STGNRDE/torchcde/log_ode.py b/model/STGNRDE/torchcde/log_ode.py new file mode 100644 index 0000000..e192282 --- /dev/null +++ b/model/STGNRDE/torchcde/log_ode.py @@ -0,0 +1,134 @@ +try: + import signatory +except ImportError: + class DummyModule: + def __getattr__(self, item): + raise ImportError("signatory has not been installed. Please install it from " + "https://github.com/patrick-kidger/signatory to use the log-ODE method.") + signatory = DummyModule() +import torch + +from . import interpolation_linear +from . import misc + + +def _logsignature_windows(x, depth, window_length, t, _version): + t = misc.validate_input_path(x, t) + + # slightly roundabout way of doing things (rather than using arange) so that it's constructed differentiably + timespan = t[-1] - t[0] + num_pieces = (timespan / window_length).ceil().to(int).item() + end_t = t[0] + num_pieces * window_length + new_t = torch.linspace(t[0], end_t, num_pieces + 1, dtype=t.dtype, device=t.device) + new_t = torch.min(new_t, t.max()) + + t_index = 0 + new_t_unique = [] + new_t_indices = [] + for new_t_elem in new_t: + while True: + lequal = (new_t_elem <= t[t_index]) + close = new_t_elem.allclose(t[t_index]) + if lequal or close: + break + t_index += 1 + new_t_indices.append(t_index + len(new_t_unique)) + if close: + continue + new_t_unique.append(new_t_elem.unsqueeze(0)) + + batch_dimensions = x.shape[:-2] + + missing_X = torch.full((1,), float('nan'), dtype=x.dtype, device=x.device).expand(*batch_dimensions, 1, x.size(-1)) + if len(new_t_unique) > 0: # no-op if len == 0, so skip for efficiency + t, indices = torch.cat([t, *new_t_unique]).sort() + x = torch.cat([x, missing_X], dim=-2)[..., indices.clamp(0, x.size(-2)), :] + + # Fill in any missing data linearly (linearly because that's what signatures do in between observations anyway) + # and conveniently that's what this already does. Here 'missing data' includes the NaNs we've just added. + x = interpolation_linear.linear_interpolation_coeffs(x, t) + + # Flatten batch dimensions for compatibility with Signatory + # flatten_X = x.view(-1, x.size(-2), x.size(-1)) + flatten_X = x.reshape(-1, x.size(-2), x.size(-1)) + first_increment = torch.zeros(*batch_dimensions, signatory.logsignature_channels(x.size(-1), depth), dtype=x.dtype, + device=x.device) + first_increment[..., :x.size(-1)] = x[..., 0, :] + logsignatures = [first_increment] + compute_logsignature = signatory.Logsignature(depth=depth) + for index, next_index, time, next_time in zip(new_t_indices[:-1], new_t_indices[1:], new_t[:-1], new_t[1:]): + logsignature = compute_logsignature(flatten_X[..., index:next_index + 1, :]) + logsignature = logsignature.view(*batch_dimensions, -1) + if _version == 0: + logsignature = logsignature * (next_time - time) + elif _version == 1: + pass + else: + raise RuntimeError + logsignatures.append(logsignature) + + logsignatures = torch.stack(logsignatures, dim=-2) + logsignatures = logsignatures.cumsum(dim=-2) + + if _version == 0: + return logsignatures, new_t + elif _version == 1: + return logsignatures + else: + raise RuntimeError + + +def logsignature_windows(x, depth, window_length, t=None): + """Calculates logsignatures over multiple windows, for the batch of controls given, as in the log-ODE method. + + ******************** + DEPRECATED: this now exists for backward compatibility. For new projects please use `logsig_windows` instead, + which has a corrected rescaling coefficient. + ******************** + + This corresponds to a transform of the time series, and should be used prior to applying one of the interpolation + schemes. + + Arguments: + x: tensor of values, of shape (..., length, input_channels), where ... is some number of batch dimensions. This + is interpreted as a (batch of) paths taking values in an input_channels-dimensional real vector space, with + length-many observations. Missing values are supported, and should be represented as NaNs. + depth: What depth to compute the logsignatures to. + window_length: How long a time interval to compute logsignatures over. + t: Optional one dimensional tensor of times. Must be monotonically increasing. If not passed will default to + tensor([0., 1., ..., length - 1]). + + Warning: + If there are missing values then calling this function can be pretty slow. Make sure to cache the result, and + don't reinstantiate it on every forward pass, if at all possible. + + Returns: + A tuple of two tensors, which are the values and times of the transformed path. + """ + return _logsignature_windows(x, depth, window_length, t, _version=0) + + +def logsig_windows(x, depth, window_length, t=None): + """Calculates logsignatures over multiple windows, for the batch of controls given, as in the log-ODE method. + + This corresponds to a transform of the time series, and should be used prior to applying one of the interpolation + schemes. + + Arguments: + x: tensor of values, of shape (..., length, input_channels), where ... is some number of batch dimensions. This + is interpreted as a (batch of) paths taking values in an `input_channels`-dimensional real vector space, + with `length`-many observations. Missing values are supported, and should be represented as NaNs. + depth: What depth to compute the logsignatures to. + window_length: How long a time interval to compute logsignatures over. + t: Optional one dimensional tensor of times. Must be monotonically increasing. If not passed will default to + `tensor([0., 1., ..., length - 1])`. + + Warning: + If there are missing values then calling this function can be pretty slow. Make sure to cache the result, and + don't reinstantiate it on every forward pass, if at all possible. + + Returns: + A tensor, which are the values of the transformed path. Times are _not_ returned: the return value is + always scaled such that the corresponding times are just `tensor([0., 1., ..., length - 1])`. + """ + return _logsignature_windows(x, depth, window_length, t, _version=1) diff --git a/model/STGNRDE/torchcde/misc.py b/model/STGNRDE/torchcde/misc.py new file mode 100644 index 0000000..aaf6823 --- /dev/null +++ b/model/STGNRDE/torchcde/misc.py @@ -0,0 +1,166 @@ +import math +import numpy as np +import torch + + +def cheap_stack(tensors, dim): + if len(tensors) == 1: + return tensors[0].unsqueeze(dim) + else: + return torch.stack(tensors, dim=dim) + + +def tridiagonal_solve(b, A_upper, A_diagonal, A_lower): + """Solves a tridiagonal system Ax = b. + + The arguments A_upper, A_digonal, A_lower correspond to the three diagonals of A. Letting U = A_upper, D=A_digonal + and L = A_lower, and assuming for simplicity that there are no batch dimensions, then the matrix A is assumed to be + of size (k, k), with entries: + + D[0] U[0] + L[0] D[1] U[1] + L[1] D[2] U[2] 0 + L[2] D[3] U[3] + . . . + . . . + . . . + L[k - 3] D[k - 2] U[k - 2] + 0 L[k - 2] D[k - 1] U[k - 1] + L[k - 1] D[k] + + Arguments: + b: A tensor of shape (..., k), where '...' is zero or more batch dimensions + A_upper: A tensor of shape (..., k - 1). + A_diagonal: A tensor of shape (..., k). + A_lower: A tensor of shape (..., k - 1). + + Returns: + A tensor of shape (..., k), corresponding to the x solving Ax = b + + Warning: + This implementation isn't super fast. You probably want to cache the result, if possible. + """ + + # This implementation is very much written for clarity rather than speed. + + A_upper, _ = torch.broadcast_tensors(A_upper, b[..., :-1]) + A_lower, _ = torch.broadcast_tensors(A_lower, b[..., :-1]) + A_diagonal, b = torch.broadcast_tensors(A_diagonal, b) + + channels = b.size(-1) + + new_b = np.empty(channels, dtype=object) + new_A_diagonal = np.empty(channels, dtype=object) + outs = np.empty(channels, dtype=object) + + new_b[0] = b[..., 0] + new_A_diagonal[0] = A_diagonal[..., 0] + for i in range(1, channels): + w = A_lower[..., i - 1] / new_A_diagonal[i - 1] + new_A_diagonal[i] = A_diagonal[..., i] - w * A_upper[..., i - 1] + new_b[i] = b[..., i] - w * new_b[i - 1] + + outs[channels - 1] = new_b[channels - 1] / new_A_diagonal[channels - 1] + for i in range(channels - 2, -1, -1): + outs[i] = (new_b[i] - A_upper[..., i] * outs[i + 1]) / new_A_diagonal[i] + + return torch.stack(outs.tolist(), dim=-1) + + +def validate_input_path(x, t): + if not x.is_floating_point(): + raise ValueError("X must both be floating point.") + + if x.ndimension() < 2: + raise ValueError("X must have at least two dimensions, corresponding to time and channels. It instead has " + "shape {}.".format(tuple(x.shape))) + + if t is None: + t = torch.linspace(0, x.size(-2) - 1, x.size(-2), dtype=x.dtype, device=x.device) + + if not t.is_floating_point(): + raise ValueError("t must both be floating point.") + if len(t.shape) != 1: + raise ValueError("t must be one dimensional. It instead has shape {}.".format(tuple(t.shape))) + prev_t_i = -math.inf + for t_i in t: + if t_i <= prev_t_i: + raise ValueError("t must be monotonically increasing.") + prev_t_i = t_i + + if x.size(-2) != t.size(0): + raise ValueError("The time dimension of X must equal the length of t. X has shape {} and t has shape {}, " + "corresponding to time dimensions of {} and {} respectively." + .format(tuple(x.shape), tuple(t.shape), x.size(-2), t.size(0))) + + if t.size(0) < 2: + raise ValueError("Must have a time dimension of size at least 2. It instead has shape {}, corresponding to a " + "time dimension of size {}.".format(tuple(t.shape), t.size(0))) + + return t + + +def forward_fill(x, fill_index=-2): + """Forward fills data in a torch tensor of shape (..., length, input_channels) along the length dim. + + Arguments: + x: tensor of values with first channel index being time, of shape (..., length, input_channels), where ... is + some number of batch dimensions. + fill_index: int that denotes the index to fill down. Default is -2 as we tend to use the convention (..., + length, input_channels) filling down the length dimension. + + Returns: + A tensor with forward filled data. + """ + # Checks + assert isinstance(x, torch.Tensor) + assert x.dim() >= 2 + + mask = torch.isnan(x) + if mask.any(): + cumsum_mask = (~mask).cumsum(dim=fill_index) + cumsum_mask[mask] = 0 + _, index = cumsum_mask.cummax(dim=fill_index) + x = x.gather(dim=fill_index, index=index) + + return x + + +class TupleControl(torch.nn.Module): + def __init__(self, *controls): + super(TupleControl, self).__init__() + + if len(controls) == 0: + raise ValueError("Expected one or more controls to batch together.") + + self._interval = controls[0].interval + grid_points = controls[0].grid_points + same_grid_points = True + for control in controls[1:]: + if (control.interval != self._interval).any(): + raise ValueError("Can only batch togehter controls over the same interval.") + if same_grid_points and (control.grid_points != grid_points).any(): + same_grid_points = False + + if same_grid_points: + self._grid_points = grid_points + else: + self._grid_points = None + + self.controls = torch.nn.ModuleList(controls) + + @property + def interval(self): + return self._interval + + @property + def grid_points(self): + if self._grid_points is None: + raise RuntimeError("Batch of controls have different grid points.") + return self._grid_points + + def evaluate(self, t): + return tuple(control.evaluate(t) for control in self.controls) + + def derivative(self, t): + return tuple(control.derivative(t) for control in self.controls) diff --git a/model/STGNRDE/torchcde/solver.py b/model/STGNRDE/torchcde/solver.py new file mode 100644 index 0000000..211cfcf --- /dev/null +++ b/model/STGNRDE/torchcde/solver.py @@ -0,0 +1,409 @@ +import torch +import torchdiffeq +import torchsde +import warnings + + +def _check_compatability_per_tensor_base(control_gradient, z0): + if control_gradient.shape[:-1] != z0.shape[:-1]: + raise ValueError("X.derivative did not return a tensor with the same number of batch dimensions as z0. " + "X.derivative returned shape {} (meaning {} batch dimensions), whilst z0 has shape {} " + "(meaning {} batch dimensions)." + "".format(tuple(control_gradient.shape), tuple(control_gradient.shape[:-1]), tuple(z0.shape), + tuple(z0.shape[:-1]))) + + +def _check_compatability_per_tensor_forward(control_gradient, system, z0): + _check_compatability_per_tensor_base(control_gradient, z0) + if system.shape[:-2] != z0.shape[:-1]: + raise ValueError("func did not return a tensor with the same number of batch dimensions as z0. func returned " + "shape {} (meaning {} batch dimensions), whilst z0 has shape {} (meaning {} batch" + " dimensions)." + "".format(tuple(system.shape), tuple(system.shape[:-2]), tuple(z0.shape), + tuple(z0.shape[:-1]))) + if system.size(-2) != z0.size(-1): + raise ValueError("func did not return a tensor with the same number of hidden channels as z0. func returned " + "shape {} (meaning {} channels), whilst z0 has shape {} (meaning {} channels)." + "".format(tuple(system.shape), system.size(-2), tuple(z0.shape), z0.size(-1))) + if system.size(-1) != control_gradient.size(-1): + raise ValueError("func did not return a tensor with the same number of input channels as X.derivative " + "returned. func returned shape {} (meaning {} channels), whilst X.derivative returned shape " + "{} (meaning {} channels)." + "".format(tuple(system.shape), system.size(-1), tuple(control_gradient.shape), + control_gradient.size(-1))) + + +def _check_compatability_per_tensor_prod(control_gradient, vector_field, z0): + _check_compatability_per_tensor_base(control_gradient, z0) + if vector_field.shape != z0.shape: + raise ValueError("func.prod did not return a tensor with the same shape as z0. func.prod returned shape {} " + "whilst z0 has shape {}." + "".format(tuple(vector_field.shape), tuple(z0.shape))) + + +def _check_compatability(X, func, z0, t): + if not hasattr(X, 'derivative'): + raise ValueError("X must have a 'derivative' method.") + control_gradient = X.derivative(t[0].detach()) + if hasattr(func, 'prod'): + is_prod = True + vector_field = func.prod(t[0], z0, control_gradient) + else: + is_prod = False + system = func(t[0], z0) + + if isinstance(z0, torch.Tensor): + is_tensor = True + if not isinstance(control_gradient, torch.Tensor): + raise ValueError("z0 is a tensor and so X.derivative must return a tensor as well.") + if is_prod: + if not isinstance(vector_field, torch.Tensor): + raise ValueError("z0 is a tensor and so func.prod must return a tensor as well.") + _check_compatability_per_tensor_prod(control_gradient, vector_field, z0) + else: + if not isinstance(system, torch.Tensor): + raise ValueError("z0 is a tensor and so func must return a tensor as well.") + _check_compatability_per_tensor_forward(control_gradient, system, z0) + + elif isinstance(z0, (tuple, list)): + is_tensor = False + if not isinstance(control_gradient, (tuple, list)): + raise ValueError("z0 is a tuple/list and so X.derivative must return a tuple/list as well.") + if len(z0) != len(control_gradient): + raise ValueError("z0 and X.derivative(t) must be tuples of the same length.") + if is_prod: + if not isinstance(vector_field, (tuple, list)): + raise ValueError("z0 is a tuple/list and so func.prod must return a tuple/list as well.") + if len(z0) != len(vector_field): + raise ValueError("z0 and func.prod(t, z, dXdt) must be tuples of the same length.") + for control_gradient_, vector_Field_, z0_ in zip(control_gradient, vector_field, z0): + if not isinstance(control_gradient_, torch.Tensor): + raise ValueError("X.derivative must return a tensor or tuple of tensors.") + if not isinstance(vector_Field_, torch.Tensor): + raise ValueError("func.prod must return a tensor or tuple/list of tensors.") + _check_compatability_per_tensor_prod(control_gradient_, vector_Field_, z0_) + else: + if not isinstance(system, (tuple, list)): + raise ValueError("z0 is a tuple/list and so func must return a tuple/list as well.") + if len(z0) != len(system): + raise ValueError("z0 and func(t, z) must be tuples of the same length.") + for control_gradient_, system_, z0_ in zip(control_gradient, system, z0): + if not isinstance(control_gradient_, torch.Tensor): + raise ValueError("X.derivative must return a tensor or tuple of tensors.") + if not isinstance(system_, torch.Tensor): + raise ValueError("func must return a tensor or tuple/list of tensors.") + _check_compatability_per_tensor_forward(control_gradient_, system_, z0_) + + else: + raise ValueError("z0 must either a tensor or a tuple/list of tensors.") + + return is_tensor, is_prod + + +class _VectorField(torch.nn.Module): + def __init__(self, X, func, is_tensor, is_prod): + super(_VectorField, self).__init__() + + self.X = X + self.func = func + self.is_tensor = is_tensor + self.is_prod = is_prod + + # torchsde backend + self.sde_type = getattr(func, "sde_type", "stratonovich") + self.noise_type = getattr(func, "noise_type", "additive") + + # torchdiffeq backend + def forward(self, t, z): + # control_gradient is of shape (..., input_channels) + control_gradient = self.X.derivative(t) + + if self.is_prod: + # out is of shape (..., hidden_channels) + out = self.func.prod(t, z, control_gradient) + else: + # vector_field is of shape (..., hidden_channels, input_channels) + vector_field = self.func(t, z) + if self.is_tensor: + # out is of shape (..., hidden_channels) + # (The squeezing is necessary to make the matrix-multiply properly batch in all cases) + out = (vector_field @ control_gradient.unsqueeze(-1)).squeeze(-1) + else: + out = tuple((vector_field_ @ control_gradient_.unsqueeze(-1)).squeeze(-1) + for vector_field_, control_gradient_ in zip(vector_field, control_gradient)) + + return out + + # torchsde backend + f = forward + + def g(self, t, z): + return torch.zeros_like(z).unsqueeze(-1) + + +def cdeint(X, func, z0, t, adjoint=True, backend="torchdiffeq", **kwargs): + r"""Solves a system of controlled differential equations. + + Solves the controlled problem: + ``` + z_t = z_{t_0} + \int_{t_0}^t f(s, z_s) dX_s + ``` + where z is a tensor of any shape, and X is some controlling signal. + + Arguments: + X: The control. This should be a instance of `torch.nn.Module`, with a `derivative` method. For example + `torchcde.CubicSpline`. This represents a continuous path derived from the data. The + derivative at a point will be computed via `X.derivative(t)`, where t is a scalar tensor. The returned + tensor should have shape (..., input_channels), where '...' is some number of batch dimensions and + input_channels is the number of channels in the input path. + func: Should be a callable describing the vector field f(t, z). If using `adjoint=True` (the default), then + should be an instance of `torch.nn.Module`, to collect the parameters for the adjoint pass. Will be called + with a scalar tensor t and a tensor z of shape (..., hidden_channels), and should return a tensor of shape + (..., hidden_channels, input_channels), where hidden_channels and input_channels are integers defined by the + `hidden_shape` and `X` arguments as above. The '...' corresponds to some number of batch dimensions. If it + has a method `prod` then that will be called to calculate the matrix-vector product f(t, z) dX_t/dt, via + `func.prod(t, z, dXdt)`. + z0: The initial state of the solution. It should have shape (..., hidden_channels), where '...' is some number + of batch dimensions. + t: a one dimensional tensor describing the times to range of times to integrate over and output the results at. + The initial time will be t[0] and the final time will be t[-1]. + adjoint: A boolean; whether to use the adjoint method to backpropagate. Defaults to True. + backend: Either "torchdiffeq" or "torchsde". Which library to use for the solvers. Note that if using torchsde + that the Brownian motion component is completely ignored -- so it's still reducing the CDE to an ODE -- + but it makes it possible to e.g. use an SDE solver there as the ODE/CDE solver here, if for some reason + that's desired. + **kwargs: Any additional kwargs to pass to the odeint solver of torchdiffeq (the most common are `rtol`, `atol`, + `method`, `options`) or the sdeint solver of torchsde. + + Returns: + The value of each z_{t_i} of the solution to the CDE z_t = z_{t_0} + \int_0^t f(s, z_s)dX_s, where t_i = t[i]. + This will be a tensor of shape (..., len(t), hidden_channels). + + Raises: + ValueError for malformed inputs. + + Note: + Supports tupled input, i.e. z0 can be a tuple of tensors, and X.derivative and func can return tuples of tensors + of the same length. + + Warnings: + Note that the returned tensor puts the sequence dimension second-to-last, rather than first like in + `torchdiffeq.odeint` or `torchsde.sdeint`. + """ + + # Reduce the default values for the tolerances because CDEs are difficult to solve with the default high tolerances. + if 'atol' not in kwargs: + kwargs['atol'] = 1e-6 + if 'rtol' not in kwargs: + kwargs['rtol'] = 1e-4 + if adjoint: + if "adjoint_atol" not in kwargs: + kwargs["adjoint_atol"] = kwargs["atol"] + if "adjoint_rtol" not in kwargs: + kwargs["adjoint_rtol"] = kwargs["rtol"] + + is_tensor, is_prod = _check_compatability(X, func, z0, t) + + if adjoint and 'adjoint_params' not in kwargs: + for buffer in X.buffers(): + # Compare based on id to avoid PyTorch not playing well with using `in` on tensors. + if buffer.requires_grad: + warnings.warn("One of the inputs to the control path X requires gradients but " + "`kwargs['adjoint_params']` has not been passed. This is probably a mistake: these " + "inputs will not receive a gradient when using the adjoint method. Either have the input " + "not require gradients (if that was unintended), or include it (and every other " + "parameter needing gradients) in `adjoint_params`. For example:\n" + "```\n" + "coeffs = ...\n" + "func = ...\n" + "X = CubicSpline(coeffs)\n" + "adjoint_params = tuple(func.parameters()) + (coeffs,)\n" + "cdeint(X=X, func=func, ..., adjoint_params=adjoint_params)\n" + "```") + + vector_field = _VectorField(X=X, func=func, is_tensor=is_tensor, is_prod=is_prod) + if backend == "torchdiffeq": + odeint = torchdiffeq.odeint_adjoint if adjoint else torchdiffeq.odeint + out = odeint(func=vector_field, y0=z0, t=t, **kwargs) + elif backend == "torchsde": + sdeint = torchsde.sdeint_adjoint if adjoint else torchsde.sdeint + out = sdeint(sde=vector_field, y0=z0, ts=t, **kwargs) + else: + raise ValueError(f"Unrecognised backend={backend}") + + if is_tensor: + batch_dims = range(1, len(out.shape) - 1) + out = out.permute(*batch_dims, 0, -1) + else: + out_ = [] + for outi in out: + batch_dims = range(1, len(outi.shape) - 1) + outi = outi.permute(*batch_dims, 0, -1) + out_.append(outi) + out = tuple(out_) + + return out + +class _VectorFieldCustom(torch.nn.Module): + def __init__(self, X, func_f, func_g, is_tensor, is_prod): + super(_VectorFieldCustom, self).__init__() + + self.X = X + self.func_f = func_f + self.func_g = func_g + self.is_tensor = is_tensor + self.is_prod = is_prod + + # torchsde backend + self.sde_type = getattr(func_f, "sde_type", "stratonovich") + self.noise_type = getattr(func_f, "noise_type", "additive") + + self.sde_type = getattr(func_g, "sde_type", "stratonovich") + self.noise_type = getattr(func_g, "noise_type", "additive") + + # torchdiffeq backend + def forward(self, t, z): + # control_gradient is of shape (..., input_channels) + control_gradient = self.X.derivative(t) + h = z[0] + z = z[1] + + if self.is_prod: + # out is of shape (..., hidden_channels) + # FIXME: i don't know it is correct + out_f = self.func_f.prod(t, h, control_gradient) + out_g = self.func_g.prod(t, z, control_gradient) + else: + # vector_field is of shape (..., hidden_channels, input_channels) + vector_field_f = self.func_f(t, h) + vector_field_g = self.func_g(t, z) + if self.is_tensor: + # out is of shape (..., hidden_channels) + # (The squeezing is necessary to make the matrix-multiply properly batch in all cases) + dh = (vector_field_f @ control_gradient.unsqueeze(-1)).squeeze(-1) + vector_field_gf = vector_field_g @ vector_field_f + out = (vector_field_gf @ control_gradient.unsqueeze(-1)).squeeze(-1) + else: + dh = tuple((vector_field_f_ @ control_gradient_.unsqueeze(-1)).squeeze(-1) + for vector_field_f_, control_gradient_ in zip(vector_field_f, control_gradient)) + vector_field_gf = vector_field_g @ vector_field_f + out = tuple((vector_field_gf_ @ control_gradient_.unsqueeze(-1)).squeeze(-1) + for vector_field_gf_, control_gradient_ in zip(vector_field_gf, control_gradient)) + # FIXME: return value to tuple + # import pdb; pdb.set_trace() + return tuple([dh, out]) + + # torchsde backend + f = forward + + def g(self, t, z): + return torch.zeros_like(z).unsqueeze(-1) + +def cdeint_custom(X, func_f, func_g, h0, z0, t, adjoint=True, backend="torchdiffeq", **kwargs): + r"""Solves a system of controlled differential equations. + + Solves the controlled problem: + ``` + z_t = z_{t_0} + \int_{t_0}^t f(s, z_s) dX_s + ``` + where z is a tensor of any shape, and X is some controlling signal. + + Arguments: + X: The control. This should be a instance of `torch.nn.Module`, with a `derivative` method. For example + `torchcde.CubicSpline`. This represents a continuous path derived from the data. The + derivative at a point will be computed via `X.derivative(t)`, where t is a scalar tensor. The returned + tensor should have shape (..., input_channels), where '...' is some number of batch dimensions and + input_channels is the number of channels in the input path. + func: Should be a callable describing the vector field f(t, z). If using `adjoint=True` (the default), then + should be an instance of `torch.nn.Module`, to collect the parameters for the adjoint pass. Will be called + with a scalar tensor t and a tensor z of shape (..., hidden_channels), and should return a tensor of shape + (..., hidden_channels, input_channels), where hidden_channels and input_channels are integers defined by the + `hidden_shape` and `X` arguments as above. The '...' corresponds to some number of batch dimensions. If it + has a method `prod` then that will be called to calculate the matrix-vector product f(t, z) dX_t/dt, via + `func.prod(t, z, dXdt)`. + z0: The initial state of the solution. It should have shape (..., hidden_channels), where '...' is some number + of batch dimensions. + t: a one dimensional tensor describing the times to range of times to integrate over and output the results at. + The initial time will be t[0] and the final time will be t[-1]. + adjoint: A boolean; whether to use the adjoint method to backpropagate. Defaults to True. + backend: Either "torchdiffeq" or "torchsde". Which library to use for the solvers. Note that if using torchsde + that the Brownian motion component is completely ignored -- so it's still reducing the CDE to an ODE -- + but it makes it possible to e.g. use an SDE solver there as the ODE/CDE solver here, if for some reason + that's desired. + **kwargs: Any additional kwargs to pass to the odeint solver of torchdiffeq (the most common are `rtol`, `atol`, + `method`, `options`) or the sdeint solver of torchsde. + + Returns: + The value of each z_{t_i} of the solution to the CDE z_t = z_{t_0} + \int_0^t f(s, z_s)dX_s, where t_i = t[i]. + This will be a tensor of shape (..., len(t), hidden_channels). + + Raises: + ValueError for malformed inputs. + + Note: + Supports tupled input, i.e. z0 can be a tuple of tensors, and X.derivative and func can return tuples of tensors + of the same length. + + Warnings: + Note that the returned tensor puts the sequence dimension second-to-last, rather than first like in + `torchdiffeq.odeint` or `torchsde.sdeint`. + """ + + # Reduce the default values for the tolerances because CDEs are difficult to solve with the default high tolerances. + if 'atol' not in kwargs: + kwargs['atol'] = 1e-6 + if 'rtol' not in kwargs: + kwargs['rtol'] = 1e-4 + if adjoint: + if "adjoint_atol" not in kwargs: + kwargs["adjoint_atol"] = kwargs["atol"] + if "adjoint_rtol" not in kwargs: + kwargs["adjoint_rtol"] = kwargs["rtol"] + + is_tensor, is_prod = _check_compatability(X, func_f, h0, t) + # is_tensor, is_prod = _check_compatability(X, func_g, z0, t) + + + if adjoint and 'adjoint_params' not in kwargs: + for buffer in X.buffers(): + # Compare based on id to avoid PyTorch not playing well with using `in` on tensors. + if buffer.requires_grad: + warnings.warn("One of the inputs to the control path X requires gradients but " + "`kwargs['adjoint_params']` has not been passed. This is probably a mistake: these " + "inputs will not receive a gradient when using the adjoint method. Either have the input " + "not require gradients (if that was unintended), or include it (and every other " + "parameter needing gradients) in `adjoint_params`. For example:\n" + "```\n" + "coeffs = ...\n" + "func = ...\n" + "X = CubicSpline(coeffs)\n" + "adjoint_params = tuple(func.parameters()) + (coeffs,)\n" + "cdeint(X=X, func=func, ..., adjoint_params=adjoint_params)\n" + "```") + + vector_field = _VectorFieldCustom(X=X, func_f=func_f, func_g=func_g, is_tensor=is_tensor, is_prod=is_prod) + if backend == "torchdiffeq": + # import pdb; pdb.set_trace() + z0 = (h0,z0) + odeint = torchdiffeq.odeint_adjoint if adjoint else torchdiffeq.odeint + out = odeint(func=vector_field, y0=z0, t=t, **kwargs) + elif backend == "torchsde": + sdeint = torchsde.sdeint_adjoint if adjoint else torchsde.sdeint + out = sdeint(sde=vector_field, y0=z0, ts=t, **kwargs) + else: + raise ValueError(f"Unrecognised backend={backend}") + + if is_tensor: + # import pdb; pdb.set_trace() + out = out[-1] + # batch_dims = range(1, len(out[-1].shape) - 1) + batch_dims = range(1, len(out.shape) - 1) + out = out.permute(*batch_dims, 0, -1) + else: + out_ = [] + for outi in out: + batch_dims = range(1, len(outi.shape) - 1) + outi = outi.permute(*batch_dims, 0, -1) + out_.append(outi) + out = tuple(out_) + return out \ No newline at end of file diff --git a/model/STGNRDE/vector_fields.py b/model/STGNRDE/vector_fields.py new file mode 100755 index 0000000..03f1434 --- /dev/null +++ b/model/STGNRDE/vector_fields.py @@ -0,0 +1,341 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class FinalTanh_f(nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers): + super(FinalTanh_f, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + self.linear_in = nn.Linear(hidden_channels, hidden_hidden_channels) + + self.linears = nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + for _ in range(num_hidden_layers - 1)) + self.linear_out = nn.Linear(hidden_hidden_channels, input_channels * hidden_channels) #32,32*4 -> # 32,32,4 + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, *args): + z = args[0] if len(args) == 1 else args[1] + z = self.linear_in(z) + z = z.relu() + + for linear in self.linears: + z = linear(z) + z = z.relu() + # z: torch.Size([64, 207, 32]) + # self.linear_out(z): torch.Size([64, 207, 64]) + z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.input_channels) + z = z.tanh() + return z + +class FinalTanh_f_prime(nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers): + super(FinalTanh_f_prime, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + self.linear_in = nn.Linear(hidden_channels, hidden_hidden_channels) + + self.linears = nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + for _ in range(num_hidden_layers - 1)) + # self.linear_out = nn.Linear(hidden_hidden_channels, input_channels * hidden_channels) #32,32*4 -> # 32,32,4 + self.linear_out = nn.Linear(hidden_hidden_channels, hidden_channels * hidden_channels) #32,32*4 -> # 32,32,4 + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, *args): + z = args[0] if len(args) == 1 else args[1] + z = self.linear_in(z) + z = z.relu() + + for linear in self.linears: + z = linear(z) + z = z.relu() + # z: torch.Size([64, 207, 32]) + # self.linear_out(z): torch.Size([64, 207, 64]) + # z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.input_channels) + z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.hidden_channels) + z = z.tanh() + return z + +class FinalTanh_f2(torch.nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers): + super(FinalTanh_f2, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + # self.linear_in = torch.nn.Linear(hidden_channels, hidden_hidden_channels) + + # self.linears = torch.nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + # for _ in range(num_hidden_layers - 1)) + # self.linear_out = torch.nn.Linear(hidden_hidden_channels, input_channels * hidden_channels) #32,32*4 -> # 32,32,4 + + self.start_conv = torch.nn.Conv2d(in_channels=hidden_channels, + out_channels=hidden_channels, + kernel_size=(1,1)) + + # self.linear = torch.nn.Conv2d(in_channels=hidden_channels, + # out_channels=hidden_channels, + # kernel_size=(1,1)) + + self.linears = torch.nn.ModuleList(torch.nn.Conv2d(in_channels=hidden_channels, + out_channels=hidden_channels, + kernel_size=(1,1)) + for _ in range(num_hidden_layers - 1)) + + self.linear_out = torch.nn.Conv2d(in_channels=hidden_channels, + out_channels=input_channels*hidden_channels, + kernel_size=(1,1)) + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, *args): + # z: torch.Size([64, 207, 32]) + z = args[0] if len(args) == 1 else args[1] + z = self.start_conv(z.transpose(1,2).unsqueeze(-1)) + z = z.relu() + + for linear in self.linears: + z = linear(z) + z = z.relu() + + z = self.linear_out(z).squeeze().transpose(1,2).view(*z.transpose(1,2).shape[:-2], self.hidden_channels, self.input_channels) + z = z.tanh() + return z + +class VectorField_g(torch.nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers, num_nodes, cheb_k, embed_dim, + g_type): + super(VectorField_g, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + self.linear_in = torch.nn.Linear(hidden_channels, hidden_hidden_channels) + + # self.linears = torch.nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + # for _ in range(num_hidden_layers - 1)) + + # project to (hidden_channels, input_channels) for cdeint requirement + self.linear_out = torch.nn.Linear(hidden_hidden_channels, hidden_channels * input_channels) + + self.g_type = g_type + if self.g_type == 'agc': + self.node_embeddings = nn.Parameter(torch.randn(num_nodes, embed_dim), requires_grad=True) + self.cheb_k = cheb_k + self.weights_pool = nn.Parameter(torch.FloatTensor(embed_dim, cheb_k, hidden_hidden_channels, hidden_hidden_channels)) + self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, hidden_hidden_channels)) + + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, *args): + z = args[0] if len(args) == 1 else args[1] + z = self.linear_in(z) + z = z.relu() + + if self.g_type == 'agc': + z = self.agc(z) + else: + raise ValueError('Check g_type argument') + # for linear in self.linears: + # z = linear(x_gconv) + # z = z.relu() + + # output shape (..., hidden_channels, input_channels) + z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.input_channels) + z = z.tanh() + return z #torch.Size([64, 307, 64, 1]) + + def agc(self, z): + """ + Adaptive Graph Convolution + - Node Adaptive Parameter Learning + - Data Adaptive Graph Generation + """ + node_num = self.node_embeddings.shape[0] + supports = F.softmax(F.relu(torch.mm(self.node_embeddings, self.node_embeddings.transpose(0, 1))), dim=1) + # laplacian=False + laplacian=False + if laplacian == True: + # support_set = [torch.eye(node_num).to(supports.device), -supports] + support_set = [supports, -torch.eye(node_num).to(supports.device)] + # support_set = [torch.eye(node_num).to(supports.device), -supports] + # support_set = [-supports] + else: + support_set = [torch.eye(node_num).to(supports.device), supports] + #default cheb_k = 3 + for k in range(2, self.cheb_k): + support_set.append(torch.matmul(2 * supports, support_set[-1]) - support_set[-2]) + supports = torch.stack(support_set, dim=0) + weights = torch.einsum('nd,dkio->nkio', self.node_embeddings, self.weights_pool) #N, cheb_k, dim_in, dim_out + bias = torch.matmul(self.node_embeddings, self.bias_pool) #N, dim_out + x_g = torch.einsum("knm,bmc->bknc", supports, z) #B, cheb_k, N, dim_in + x_g = x_g.permute(0, 2, 1, 3) # B, N, cheb_k, dim_in + z = torch.einsum('bnki,nkio->bno', x_g, weights) + bias #b, N, dim_out + return z + + +class VectorField_only_g(torch.nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers, num_nodes, cheb_k, embed_dim, + g_type): + super(VectorField_only_g, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + self.linear_in = torch.nn.Linear(hidden_channels, hidden_hidden_channels) + + # self.linears = torch.nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + # for _ in range(num_hidden_layers - 1)) + + #FIXME: + self.linear_out = torch.nn.Linear(hidden_hidden_channels, input_channels * hidden_channels) #32,32*4 -> # 32,32,4 + # self.linear_out = torch.nn.Linear(hidden_hidden_channels, hidden_channels * hidden_channels) #32,32*4 -> # 32,32,4 + + self.g_type = g_type + if self.g_type == 'agc': + self.node_embeddings = nn.Parameter(torch.randn(num_nodes, embed_dim), requires_grad=True) + self.cheb_k = cheb_k + self.weights_pool = nn.Parameter(torch.FloatTensor(embed_dim, cheb_k, hidden_hidden_channels, hidden_hidden_channels)) + self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, hidden_hidden_channels)) + + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, *args): + z = args[0] if len(args) == 1 else args[1] + z = self.linear_in(z) + z = z.relu() + + if self.g_type == 'agc': + z = self.agc(z) + else: + raise ValueError('Check g_type argument') + # for linear in self.linears: + # z = linear(x_gconv) + # z = z.relu() + + #FIXME: + z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.input_channels) + # z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.hidden_channels) + z = z.tanh() + return z #torch.Size([64, 307, 64, 1]) + + def agc(self, z): + """ + Adaptive Graph Convolution + - Node Adaptive Parameter Learning + - Data Adaptive Graph Generation + """ + node_num = self.node_embeddings.shape[0] + supports = F.softmax(F.relu(torch.mm(self.node_embeddings, self.node_embeddings.transpose(0, 1))), dim=1) + + laplacian=False + if laplacian == True: + # support_set = [torch.eye(node_num).to(supports.device), -supports] + support_set = [supports, -torch.eye(node_num).to(supports.device)] + # support_set = [torch.eye(node_num).to(supports.device), -supports] + # support_set = [-supports] + else: + support_set = [torch.eye(node_num).to(supports.device), supports] + #default cheb_k = 3 + for k in range(2, self.cheb_k): + support_set.append(torch.matmul(2 * supports, support_set[-1]) - support_set[-2]) + supports = torch.stack(support_set, dim=0) + weights = torch.einsum('nd,dkio->nkio', self.node_embeddings, self.weights_pool) #N, cheb_k, dim_in, dim_out + bias = torch.matmul(self.node_embeddings, self.bias_pool) #N, dim_out + x_g = torch.einsum("knm,bmc->bknc", supports, z) #B, cheb_k, N, dim_in + x_g = x_g.permute(0, 2, 1, 3) # B, N, cheb_k, dim_in + z = torch.einsum('bnki,nkio->bno', x_g, weights) + bias #b, N, dim_out + return z + +class VectorField_g_prime(torch.nn.Module): + def __init__(self, input_channels, hidden_channels, hidden_hidden_channels, num_hidden_layers, num_nodes, cheb_k, embed_dim, + g_type): + super(VectorField_g_prime, self).__init__() + + self.input_channels = input_channels + self.hidden_channels = hidden_channels + self.hidden_hidden_channels = hidden_hidden_channels + self.num_hidden_layers = num_hidden_layers + + self.linear_in = torch.nn.Linear(hidden_channels, hidden_hidden_channels) + + # self.linears = torch.nn.ModuleList(torch.nn.Linear(hidden_hidden_channels, hidden_hidden_channels) + # for _ in range(num_hidden_layers - 1)) + + self.linear_out = torch.nn.Linear(hidden_hidden_channels, input_channels * hidden_channels) #32,32*4 -> # 32,32,4 + + self.g_type = g_type + if self.g_type == 'agc': + self.node_embeddings = nn.Parameter(torch.randn(num_nodes, embed_dim), requires_grad=True) + self.cheb_k = cheb_k + self.weights_pool = nn.Parameter(torch.FloatTensor(embed_dim, cheb_k, hidden_hidden_channels, hidden_hidden_channels)) + self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, hidden_hidden_channels)) + + + def extra_repr(self): + return "input_channels: {}, hidden_channels: {}, hidden_hidden_channels: {}, num_hidden_layers: {}" \ + "".format(self.input_channels, self.hidden_channels, self.hidden_hidden_channels, self.num_hidden_layers) + + def forward(self, z): + z = self.linear_in(z) + z = z.relu() + + if self.g_type == 'agc': + z = self.agc(z) + else: + raise ValueError('Check g_type argument') + # for linear in self.linears: + # z = linear(x_gconv) + # z = z.relu() + + z = self.linear_out(z).view(*z.shape[:-1], self.hidden_channels, self.input_channels) + z = z.tanh() + return z #torch.Size([64, 307, 64, 1]) + + def agc(self, z): + """ + Adaptive Graph Convolution + - Node Adaptive Parameter Learning + - Data Adaptive Graph Generation + """ + node_num = self.node_embeddings.shape[0] + supports = F.softmax(F.relu(torch.mm(self.node_embeddings, self.node_embeddings.transpose(0, 1))), dim=1) + support_set = [torch.eye(node_num).to(supports.device), supports] + #default cheb_k = 3 + for k in range(2, self.cheb_k): + support_set.append(torch.matmul(2 * supports, support_set[-1]) - support_set[-2]) + supports = torch.stack(support_set, dim=0) + weights = torch.einsum('nd,dkio->nkio', self.node_embeddings, self.weights_pool) #N, cheb_k, dim_in, dim_out + bias = torch.matmul(self.node_embeddings, self.bias_pool) #N, dim_out + x_g = torch.einsum("knm,bmc->bknc", supports, z) #B, cheb_k, N, dim_in + x_g = x_g.permute(0, 2, 1, 3) # B, N, cheb_k, dim_in + z = torch.einsum('bnki,nkio->bno', x_g, weights) + bias #b, N, dim_out + return z \ No newline at end of file diff --git a/model/model_selector.py b/model/model_selector.py index 9cb75f8..4fba6bc 100755 --- a/model/model_selector.py +++ b/model/model_selector.py @@ -20,6 +20,8 @@ from model.STAEFormer.STAEFormer import STAEformer from model.EXP.EXP32 import EXP as EXP from model.MegaCRN.MegaCRNModel import MegaCRNModel from model.ST_SSL.ST_SSL import STSSLModel +from model.STGNRDE.Make_model import make_model as make_nrde_model +from model.STAWnet.STAWnet import STAWnet def model_selector(model): match model['type']: @@ -45,4 +47,6 @@ def model_selector(model): case 'EXP': return EXP(model) case 'MegaCRN': return MegaCRNModel(model) case 'ST_SSL': return STSSLModel(model) + case 'STGNRDE': return make_nrde_model(model) + case 'STAWnet': return STAWnet(model) diff --git a/trainer/DCRNN_Trainer.py b/trainer/DCRNN_Trainer.py index 97a8290..3e9aa56 100755 --- a/trainer/DCRNN_Trainer.py +++ b/trainer/DCRNN_Trainer.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -34,6 +35,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) def _run_epoch(self, epoch, dataloader, mode): if mode == 'train': @@ -49,12 +52,13 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, (data, target) in enumerate(dataloader): + start_time = time.time() label = target[..., :self.args['output_dim']] - # label = target[..., :self.args['output_dim']] output = self.model(data, labels=label.clone()).to(self.args['device']) if self.args['real_value']: output = self.scaler.inverse_transform(output) + label = self.scaler.inverse_transform(label) loss = self.loss(output, label) if optimizer_step and self.optimizer is not None: @@ -65,6 +69,8 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -78,6 +84,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -94,6 +102,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -126,6 +135,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): @@ -154,11 +171,11 @@ class Trainer: y_pred.append(output) y_true.append(label) - if args['real_value']: - y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0)) - else: - y_pred = torch.cat(y_pred, dim=0) + y_pred = torch.cat(y_pred, dim=0) y_true = torch.cat(y_true, dim=0) + if args['real_value']: + y_pred = scaler.inverse_transform(y_pred) + y_true = scaler.inverse_transform(y_true) for t in range(y_true.shape[1]): mae, rmse, mape = all_metrics(y_pred[:, t, ...], y_true[:, t, ...], diff --git a/trainer/E32Trainer.py b/trainer/E32Trainer.py index 1a8d062..b1bce7c 100644 --- a/trainer/E32Trainer.py +++ b/trainer/E32Trainer.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -34,6 +35,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) def _run_epoch(self, epoch, dataloader, mode): is_train = (mode == 'train') @@ -45,6 +48,7 @@ class Trainer: tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, batch in enumerate(dataloader): + start_time = time.time() # unpack the new cycle_index data, target, cycle_index = batch data = data.to(self.args['device']) @@ -72,6 +76,8 @@ class Trainer: self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() # logging @@ -86,6 +92,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -102,6 +110,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -134,6 +143,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/EXP_trainer.py b/trainer/EXP_trainer.py index 5613870..80dc6c7 100755 --- a/trainer/EXP_trainer.py +++ b/trainer/EXP_trainer.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -34,6 +35,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) def _run_epoch(self, epoch, dataloader, mode): if mode == 'train': @@ -49,6 +52,7 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, (data, target) in enumerate(dataloader): + start_time = time.time() label = target[..., :self.args['output_dim']] output = self.model(data).to(self.args['device']) @@ -64,6 +68,9 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) + total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -77,6 +84,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -93,6 +102,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -124,7 +134,14 @@ class Trainer: torch.save(best_model, self.best_path) torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") - + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/PDG2SEQ_Trainer.py b/trainer/PDG2SEQ_Trainer.py index bde4801..a6dc908 100755 --- a/trainer/PDG2SEQ_Trainer.py +++ b/trainer/PDG2SEQ_Trainer.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -35,6 +36,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) def _run_epoch(self, epoch, dataloader, mode): if mode == 'train': @@ -50,6 +53,7 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, (data, target) in enumerate(dataloader): + start_time = time.time() self.batches_seen += 1 label = target[..., :self.args['output_dim']].clone() output = self.model(data, target, self.batches_seen).to(self.args['device']) @@ -66,6 +70,9 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + # record step time + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -79,6 +86,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -95,6 +104,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -127,6 +137,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/STMLP_Trainer.py b/trainer/STMLP_Trainer.py index 6489221..6e416e8 100644 --- a/trainer/STMLP_Trainer.py +++ b/trainer/STMLP_Trainer.py @@ -11,6 +11,7 @@ from tqdm import tqdm from lib.logger import get_logger from lib.loss_function import all_metrics from model.STMLP.STMLP import STMLP +from lib.training_stats import TrainingStats class Trainer: @@ -43,6 +44,8 @@ class Trainer: os.makedirs(self.pretrain_dir, exist_ok=True) self.logger = get_logger(self.args['log_dir'], name=self.model.__class__.__name__, debug=self.args['debug']) self.logger.info(f"Experiment log path in: {self.args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) if self.args['teacher_stu']: self.tmodel = self.loadTeacher(args) @@ -67,6 +70,7 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, (data, target) in enumerate(dataloader): + start_time = time.time() if self.args['teacher_stu']: label = target[..., :self.args['output_dim']] output, out_, _ = self.model(data) @@ -100,6 +104,8 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -113,6 +119,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -129,6 +137,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -165,6 +174,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/Trainer.py b/trainer/Trainer.py index 9fbe3a9..013d852 100755 --- a/trainer/Trainer.py +++ b/trainer/Trainer.py @@ -211,6 +211,13 @@ class Trainer: self._finalize_training(best_model, best_test_model) + # 输出参数量 + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass + def _finalize_training(self, best_model, best_test_model): self.model.load_state_dict(best_model) self.logger.info("Testing on best validation model") diff --git a/trainer/Trainer_old.py b/trainer/Trainer_old.py index 5613870..4004f8a 100755 --- a/trainer/Trainer_old.py +++ b/trainer/Trainer_old.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -34,6 +35,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) def _run_epoch(self, epoch, dataloader, mode): if mode == 'train': @@ -49,6 +52,7 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, (data, target) in enumerate(dataloader): + start_time = time.time() label = target[..., :self.args['output_dim']] output = self.model(data).to(self.args['device']) @@ -64,6 +68,8 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -77,6 +83,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): @@ -93,6 +101,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -125,6 +134,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/cdeTrainer/cdetrainer.py b/trainer/cdeTrainer/cdetrainer.py index 59fefa8..f939d66 100755 --- a/trainer/cdeTrainer/cdetrainer.py +++ b/trainer/cdeTrainer/cdetrainer.py @@ -7,6 +7,7 @@ from tqdm import tqdm import torch from lib.logger import get_logger from lib.loss_function import all_metrics +from lib.training_stats import TrainingStats class Trainer: @@ -35,6 +36,8 @@ class Trainer: os.makedirs(args['log_dir'], exist_ok=True) self.logger = get_logger(args['log_dir'], name=self.model.__class__.__name__, debug=args['debug']) self.logger.info(f"Experiment log path in: {args['log_dir']}") + # Stats tracker + self.stats = TrainingStats(device=args['device']) self.times = times.to(self.device, dtype=torch.float) self.w = w @@ -52,6 +55,7 @@ class Trainer: with torch.set_grad_enabled(optimizer_step): with tqdm(total=len(dataloader), desc=f'{mode.capitalize()} Epoch {epoch}') as pbar: for batch_idx, batch in enumerate(dataloader): + start_time = time.time() batch = tuple(b.to(self.device, dtype=torch.float) for b in batch) *train_coeffs, target = batch label = target[..., :self.args['output_dim']] @@ -69,6 +73,8 @@ class Trainer: torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['max_grad_norm']) self.optimizer.step() + step_time = time.time() - start_time + self.stats.record_step_time(step_time, mode) total_loss += loss.item() if mode == 'train' and (batch_idx + 1) % self.args['log_step'] == 0: @@ -82,6 +88,8 @@ class Trainer: avg_loss = total_loss / len(dataloader) self.logger.info( f'{mode.capitalize()} Epoch {epoch}: average Loss: {avg_loss:.6f}, time: {time.time() - epoch_time:.2f} s') + # 记录内存 + self.stats.record_memory_usage() return avg_loss def train_epoch(self, epoch): return self._run_epoch(epoch, self.train_loader, 'train') @@ -97,6 +105,7 @@ class Trainer: best_loss, best_test_loss = float('inf'), float('inf') not_improved_count = 0 + self.stats.start_training() self.logger.info("Training process started") for epoch in range(1, self.args['epochs'] + 1): train_epoch_loss = self.train_epoch(epoch) @@ -129,6 +138,14 @@ class Trainer: torch.save(best_test_model, self.best_test_path) self.logger.info(f"Best models saved at {self.best_path} and {self.best_test_path}") + # 输出统计与参数 + self.stats.end_training() + self.stats.report(self.logger) + try: + total_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad) + self.logger.info(f"Trainable params: {total_params}") + except Exception: + pass self._finalize_training(best_model, best_test_model) def _finalize_training(self, best_model, best_test_model): diff --git a/trainer/trainer_selector.py b/trainer/trainer_selector.py index 57e48d6..5f66185 100755 --- a/trainer/trainer_selector.py +++ b/trainer/trainer_selector.py @@ -11,6 +11,8 @@ def select_trainer(model, loss, optimizer, train_loader, val_loader, test_loader match args['model']['type']: case "STGNCDE": return cdeTrainer(model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args['train'], lr_scheduler, kwargs[0], None) + case "STGNRDE": return cdeTrainer(model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args['train'], + lr_scheduler, kwargs[0], None) case 'DCRNN': return DCRNN_Trainer(model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args['train'], lr_scheduler) case 'PDG2SEQ': return PDG2SEQ_Trainer(model, loss, optimizer, train_loader, val_loader, test_loader, scaler, args['train'],