更新.gitignore以忽略.temp_repo文件，更新README.md以支持新模型STAWnet，修改DCRNN和STGNRDE配置文件，添加STAWnet和STGNRDE模型及其训练器，增强模型选择器和数据加载器以支持新模型，更新损失函数以适应新需求，添加训练统计功能以跟踪训练过程中的性能指标。
This commit is contained in:
parent
6545482a58
commit
8958cd7d95
|
|
@ -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/
|
||||
|
|
@ -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’ ... 取决于机器卡数
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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':
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
@ -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).
|
||||
********************
|
||||
"""
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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, ...],
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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'],
|
||||
|
|
|
|||
Loading…
Reference in New Issue