wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-23 13:10:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Time grad (#28)

Browse Source 
* initial uncond image gaussian diff

TODO make it work for multivariate vector
add conditioning

* remove tqdm

* initial unet

TODO convert to 1d conv

* initial time grad estimator

* initial training

* initial sampling

* added huber loss

* use SinusoidalPosEmb from wavegrad

* use time diff network

* fix reshaping

* fix missing property

* clip false

* updated api

* added padding

* added circular padding

* use linear schedule

* added more schedules

* added back cosine schedule

* Delete Solar-time-grad.ipynb

* updated estimator API

* not tuple

* renamed to EpsilonTheta

* removed

* added example notebook

* removed some output

* fix requirements

* formatting

* added more options to time-grad

* added article
This commit is contained in:
Kashif Rasul
2021-02-11 10:09:25 +01:00
committed by GitHub Enterprise
parent 512b968b13
commit 908945b422
13 changed files with 2074 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+12
View File
@@ -125,3 +125,15 @@ We have implemented the following model using this framework:
eprint = {2002.06103},
}
```
* [Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting
](https://arxiv.org/abs/2101.12072)
```tex
@article{rasul2020timegrad,
Author = {Kashif Rasul, Calvin Seward, Ingmar Schuster, Roland Vollgraf}
Title = {Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting},
Year = {2021},
archivePrefix = {arXiv},
eprint = {2101.12072},
}
```
examples/Time-Grad-Electricity.ipynb
+761
View File
File diff suppressed because one or more lines are too long
pts/model/n_beats/n_beats_ensemble.py
+1
View File
@@ -169,6 +169,7 @@ class NBEATSEnsembleEstimator(PyTorchEstimator):
**kwargs
Arguments passed down to the individual estimators.
"""
@validted()
def __init__(
self,
pts/model/simple_feedforward/simple_feedforward_estimator.py
+1
View File
@@ -86,6 +86,7 @@ class SimpleFeedForwardEstimator(PyTorchEstimator):
Number of evaluation samples per time series to increase parallelism during inference.
This is a model optimization that does not affect the accuracy (default: 100)
"""
@validated()
def __init__(
self,
pts/model/time_grad/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .time_grad_estimator import TimeGradEstimator
from .time_grad_network import TimeGradTrainingNetwork, TimeGradPredictionNetwork
from .epsilon_theta import EpsilonTheta
pts/model/time_grad/epsilon_theta.py
+136
View File
@@ -0,0 +1,136 @@
import math
import torch
from torch import nn
import torch.nn.functional as F
class DiffusionEmbedding(nn.Module):
def __init__(self, dim, proj_dim, max_steps=500):
super().__init__()
self.register_buffer(
"embedding", self._build_embedding(dim, max_steps), persistent=False
)
self.projection1 = nn.Linear(dim * 2, proj_dim)
self.projection2 = nn.Linear(proj_dim, proj_dim)
def forward(self, diffusion_step):
x = self.embedding[diffusion_step]
x = self.projection1(x)
x = F.silu(x)
x = self.projection2(x)
x = F.silu(x)
return x
def _build_embedding(self, dim, max_steps):
steps = torch.arange(max_steps).unsqueeze(1) # [T,1]
dims = torch.arange(dim).unsqueeze(0) # [1,dim]
table = steps * 10.0 ** (dims * 4.0 / dim) # [T,dim]
table = torch.cat([torch.sin(table), torch.cos(table)], dim=1)
return table
class ResidualBlock(nn.Module):
def __init__(self, hidden_size, residual_channels, dilation):
super().__init__()
self.dilated_conv = nn.Conv1d(
residual_channels,
2 * residual_channels,
3,
padding=dilation,
dilation=dilation,
padding_mode="circular",
)
self.diffusion_projection = nn.Linear(hidden_size, residual_channels)
self.conditioner_projection = nn.Conv1d(
1, 2 * residual_channels, 1, padding=2, padding_mode="circular"
)
self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1)
nn.init.kaiming_normal_(self.conditioner_projection.weight)
nn.init.kaiming_normal_(self.output_projection.weight)
def forward(self, x, conditioner, diffusion_step):
diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
conditioner = self.conditioner_projection(conditioner)
y = x + diffusion_step
y = self.dilated_conv(y) + conditioner
gate, filter = torch.chunk(y, 2, dim=1)
y = torch.sigmoid(gate) * torch.tanh(filter)
y = self.output_projection(y)
y = F.leaky_relu(y, 0.4)
residual, skip = torch.chunk(y, 2, dim=1)
return (x + residual) / math.sqrt(2.0), skip
class CondUpsampler(nn.Module):
def __init__(self, cond_length, target_dim):
super().__init__()
self.linear1 = nn.Linear(cond_length, target_dim // 2)
self.linear2 = nn.Linear(target_dim // 2, target_dim)
def forward(self, x):
x = self.linear1(x)
x = F.leaky_relu(x, 0.4)
x = self.linear2(x)
x = F.leaky_relu(x, 0.4)
return x
class EpsilonTheta(nn.Module):
def __init__(
self,
target_dim,
cond_length,
time_emb_dim=16,
residual_layers=8,
residual_channels=8,
dilation_cycle_length=2,
residual_hidden=64,
):
super().__init__()
self.input_projection = nn.Conv1d(
1, residual_channels, 1, padding=2, padding_mode="circular"
)
self.diffusion_embedding = DiffusionEmbedding(
time_emb_dim, proj_dim=residual_hidden
)
self.cond_upsampler = CondUpsampler(
target_dim=target_dim, cond_length=cond_length
)
self.residual_layers = nn.ModuleList(
[
ResidualBlock(
residual_channels=residual_channels,
dilation=2 ** (i % dilation_cycle_length),
hidden_size=residual_hidden,
)
for i in range(residual_layers)
]
)
self.skip_projection = nn.Conv1d(residual_channels, residual_channels, 3)
self.output_projection = nn.Conv1d(residual_channels, 1, 3)
nn.init.kaiming_normal_(self.input_projection.weight)
nn.init.kaiming_normal_(self.skip_projection.weight)
nn.init.zeros_(self.output_projection.weight)
def forward(self, inputs, time, cond):
x = self.input_projection(inputs)
x = F.leaky_relu(x, 0.4)
diffusion_step = self.diffusion_embedding(time)
cond_up = self.cond_upsampler(cond)
skip = []
for layer in self.residual_layers:
x, skip_connection = layer(x, cond_up, diffusion_step)
skip.append(skip_connection)
x = torch.sum(torch.stack(skip), dim=0) / math.sqrt(len(self.residual_layers))
x = self.skip_projection(x)
x = F.leaky_relu(x, 0.4)
x = self.output_projection(x)
return x
pts/model/time_grad/time_grad_estimator.py
+257
View File
@@ -0,0 +1,257 @@
from typing import List, Optional
import torch
from gluonts.dataset.field_names import FieldName
from gluonts.time_feature import TimeFeature
from gluonts.torch.model.predictor import PyTorchPredictor
from gluonts.torch.support.util import copy_parameters
from gluonts.model.predictor import Predictor
from gluonts.transform import (
Transformation,
Chain,
InstanceSplitter,
ExpectedNumInstanceSampler,
ValidationSplitSampler,
TestSplitSampler,
RenameFields,
AsNumpyArray,
ExpandDimArray,
AddObservedValuesIndicator,
AddTimeFeatures,
VstackFeatures,
SetFieldIfNotPresent,
TargetDimIndicator,
)
from pts import Trainer
from pts.feature import (
fourier_time_features_from_frequency,
lags_for_fourier_time_features_from_frequency,
)
from pts.model import PyTorchEstimator
from pts.model.utils import get_module_forward_input_names
from .time_grad_network import TimeGradTrainingNetwork, TimeGradPredictionNetwork
class TimeGradEstimator(PyTorchEstimator):
def __init__(
self,
input_size: int,
freq: str,
prediction_length: int,
target_dim: int,
trainer: Trainer = Trainer(),
context_length: Optional[int] = None,
num_layers: int = 2,
num_cells: int = 40,
cell_type: str = "LSTM",
num_parallel_samples: int = 100,
dropout_rate: float = 0.1,
cardinality: List[int] = [1],
embedding_dimension: int = 5,
conditioning_length: int = 100,
diff_steps: int = 100,
loss_type: str = "l2",
beta_end=0.1,
beta_schedule="linear",
residual_layers=8,
residual_channels=8,
dilation_cycle_length=2,
scaling: bool = True,
pick_incomplete: bool = False,
lags_seq: Optional[List[int]] = None,
time_features: Optional[List[TimeFeature]] = None,
**kwargs,
) -> None:
super().__init__(trainer=trainer, **kwargs)
self.freq = freq
self.context_length = (
context_length if context_length is not None else prediction_length
)
self.input_size = input_size
self.prediction_length = prediction_length
self.target_dim = target_dim
self.num_layers = num_layers
self.num_cells = num_cells
self.cell_type = cell_type
self.num_parallel_samples = num_parallel_samples
self.dropout_rate = dropout_rate
self.cardinality = cardinality
self.embedding_dimension = embedding_dimension
self.conditioning_length = conditioning_length
self.diff_steps = diff_steps
self.loss_type = loss_type
self.beta_end = beta_end
self.beta_schedule = beta_schedule
self.residual_layers = residual_layers
self.residual_channels = residual_channels
self.dilation_cycle_length = dilation_cycle_length
self.lags_seq = (
lags_seq
if lags_seq is not None
else lags_for_fourier_time_features_from_frequency(freq_str=freq)
)
self.time_features = (
time_features
if time_features is not None
else fourier_time_features_from_frequency(self.freq)
)
self.history_length = self.context_length + max(self.lags_seq)
self.pick_incomplete = pick_incomplete
self.scaling = scaling
self.train_sampler = ExpectedNumInstanceSampler(
num_instances=1.0,
min_past=0 if pick_incomplete else self.history_length,
min_future=prediction_length,
)
self.validation_sampler = ValidationSplitSampler(
min_past=0 if pick_incomplete else self.history_length,
min_future=prediction_length,
)
def create_transformation(self) -> Transformation:
return Chain(
[
AsNumpyArray(
field=FieldName.TARGET,
expected_ndim=2,
),
# maps the target to (1, T)
# if the target data is uni dimensional
ExpandDimArray(
field=FieldName.TARGET,
axis=None,
),
AddObservedValuesIndicator(
target_field=FieldName.TARGET,
output_field=FieldName.OBSERVED_VALUES,
),
AddTimeFeatures(
start_field=FieldName.START,
target_field=FieldName.TARGET,
output_field=FieldName.FEAT_TIME,
time_features=self.time_features,
pred_length=self.prediction_length,
),
VstackFeatures(
output_field=FieldName.FEAT_TIME,
input_fields=[FieldName.FEAT_TIME],
),
SetFieldIfNotPresent(field=FieldName.FEAT_STATIC_CAT, value=[0]),
TargetDimIndicator(
field_name="target_dimension_indicator",
target_field=FieldName.TARGET,
),
AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1),
]
)
def create_instance_splitter(self, mode: str):
assert mode in ["training", "validation", "test"]
instance_sampler = {
"training": self.train_sampler,
"validation": self.validation_sampler,
"test": TestSplitSampler(),
}[mode]
return InstanceSplitter(
target_field=FieldName.TARGET,
is_pad_field=FieldName.IS_PAD,
start_field=FieldName.START,
forecast_start_field=FieldName.FORECAST_START,
instance_sampler=instance_sampler,
past_length=self.history_length,
future_length=self.prediction_length,
time_series_fields=[
FieldName.FEAT_TIME,
FieldName.OBSERVED_VALUES,
],
) + (
RenameFields(
{
f"past_{FieldName.TARGET}": f"past_{FieldName.TARGET}_cdf",
f"future_{FieldName.TARGET}": f"future_{FieldName.TARGET}_cdf",
}
)
)
def create_training_network(self, device: torch.device) -> TimeGradTrainingNetwork:
return TimeGradTrainingNetwork(
input_size=self.input_size,
target_dim=self.target_dim,
num_layers=self.num_layers,
num_cells=self.num_cells,
cell_type=self.cell_type,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
dropout_rate=self.dropout_rate,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
diff_steps=self.diff_steps,
loss_type=self.loss_type,
beta_end=self.beta_end,
beta_schedule=self.beta_schedule,
residual_layers=self.residual_layers,
residual_channels=self.residual_channels,
dilation_cycle_length=self.dilation_cycle_length,
lags_seq=self.lags_seq,
scaling=self.scaling,
conditioning_length=self.conditioning_length,
).to(device)
def create_predictor(
self,
transformation: Transformation,
trained_network: TimeGradTrainingNetwork,
device: torch.device,
) -> Predictor:
prediction_network = TimeGradPredictionNetwork(
input_size=self.input_size,
target_dim=self.target_dim,
num_layers=self.num_layers,
num_cells=self.num_cells,
cell_type=self.cell_type,
history_length=self.history_length,
context_length=self.context_length,
prediction_length=self.prediction_length,
dropout_rate=self.dropout_rate,
cardinality=self.cardinality,
embedding_dimension=self.embedding_dimension,
diff_steps=self.diff_steps,
loss_type=self.loss_type,
beta_end=self.beta_end,
beta_schedule=self.beta_schedule,
residual_layers=self.residual_layers,
residual_channels=self.residual_channels,
dilation_cycle_length=self.dilation_cycle_length,
lags_seq=self.lags_seq,
scaling=self.scaling,
conditioning_length=self.conditioning_length,
num_parallel_samples=self.num_parallel_samples,
).to(device)
copy_parameters(trained_network, prediction_network)
input_names = get_module_forward_input_names(prediction_network)
prediction_splitter = self.create_instance_splitter("test")
return PyTorchPredictor(
input_transform=transformation + prediction_splitter,
input_names=input_names,
prediction_net=prediction_network,
batch_size=self.trainer.batch_size,
freq=self.freq,
prediction_length=self.prediction_length,
device=device,
)
pts/model/time_grad/time_grad_network.py
+604
View File
@@ -0,0 +1,604 @@
from torch.nn.modules import loss
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
from gluonts.core.component import validated
from pts.model import weighted_average
from pts.modules import GaussianDiffusion, DiffusionOutput, MeanScaler, NOPScaler
from .epsilon_theta import EpsilonTheta
class TimeGradTrainingNetwork(nn.Module):
@validated()
def __init__(
self,
input_size: int,
num_layers: int,
num_cells: int,
cell_type: str,
history_length: int,
context_length: int,
prediction_length: int,
dropout_rate: float,
lags_seq: List[int],
target_dim: int,
conditioning_length: int,
diff_steps: int,
loss_type: str,
beta_end: float,
beta_schedule: str,
residual_layers: int,
residual_channels: int,
dilation_cycle_length: int,
cardinality: List[int] = [1],
embedding_dimension: int = 1,
scaling: bool = True,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.target_dim = target_dim
self.prediction_length = prediction_length
self.context_length = context_length
self.history_length = history_length
self.scaling = scaling
assert len(set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!"
lags_seq.sort()
self.lags_seq = lags_seq
self.cell_type = cell_type
rnn_cls = {"LSTM": nn.LSTM, "GRU": nn.GRU}[cell_type]
self.rnn = rnn_cls(
input_size=input_size,
hidden_size=num_cells,
num_layers=num_layers,
dropout=dropout_rate,
batch_first=True,
)
self.denoise_fn = EpsilonTheta(
target_dim=target_dim,
cond_length=conditioning_length,
residual_layers=residual_layers,
residual_channels=residual_channels,
dilation_cycle_length=dilation_cycle_length,
)
self.diffusion = GaussianDiffusion(
self.denoise_fn,
input_size=target_dim,
diff_steps=diff_steps,
loss_type=loss_type,
beta_end=beta_end,
beta_schedule=beta_schedule,
)
self.distr_output = DiffusionOutput(
self.diffusion, input_size=target_dim, cond_size=conditioning_length
)
self.proj_dist_args = self.distr_output.get_args_proj(num_cells)
self.embed_dim = 1
self.embed = nn.Embedding(
num_embeddings=self.target_dim, embedding_dim=self.embed_dim
)
if self.scaling:
self.scaler = MeanScaler(keepdim=True)
else:
self.scaler = NOPScaler(keepdim=True)
@staticmethod
def get_lagged_subsequences(
sequence: torch.Tensor,
sequence_length: int,
indices: List[int],
subsequences_length: int = 1,
) -> torch.Tensor:
"""
Returns lagged subsequences of a given sequence.
Parameters
----------
sequence
the sequence from which lagged subsequences should be extracted.
Shape: (N, T, C).
sequence_length
length of sequence in the T (time) dimension (axis = 1).
indices
list of lag indices to be used.
subsequences_length
length of the subsequences to be extracted.
Returns
--------
lagged : Tensor
a tensor of shape (N, S, C, I),
where S = subsequences_length and I = len(indices),
containing lagged subsequences.
Specifically, lagged[i, :, j, k] = sequence[i, -indices[k]-S+j, :].
"""
# we must have: history_length + begin_index >= 0
# that is: history_length - lag_index - sequence_length >= 0
# hence the following assert
assert max(indices) + subsequences_length <= sequence_length, (
f"lags cannot go further than history length, found lag "
f"{max(indices)} while history length is only {sequence_length}"
)
assert all(lag_index >= 0 for lag_index in indices)
lagged_values = []
for lag_index in indices:
begin_index = -lag_index - subsequences_length
end_index = -lag_index if lag_index > 0 else None
lagged_values.append(sequence[:, begin_index:end_index, ...].unsqueeze(1))
return torch.cat(lagged_values, dim=1).permute(0, 2, 3, 1)
def unroll(
self,
lags: torch.Tensor,
scale: torch.Tensor,
time_feat: torch.Tensor,
target_dimension_indicator: torch.Tensor,
unroll_length: int,
begin_state: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
) -> Tuple[
torch.Tensor,
Union[List[torch.Tensor], torch.Tensor],
torch.Tensor,
torch.Tensor,
]:
# (batch_size, sub_seq_len, target_dim, num_lags)
lags_scaled = lags / scale.unsqueeze(-1)
# assert_shape(
# lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)),
# )
input_lags = lags_scaled.reshape(
(-1, unroll_length, len(self.lags_seq) * self.target_dim)
)
# (batch_size, target_dim, embed_dim)
index_embeddings = self.embed(target_dimension_indicator)
# assert_shape(index_embeddings, (-1, self.target_dim, self.embed_dim))
# (batch_size, seq_len, target_dim * embed_dim)
repeated_index_embeddings = (
index_embeddings.unsqueeze(1)
.expand(-1, unroll_length, -1, -1)
.reshape((-1, unroll_length, self.target_dim * self.embed_dim))
)
# (batch_size, sub_seq_len, input_dim)
inputs = torch.cat((input_lags, repeated_index_embeddings, time_feat), dim=-1)
# unroll encoder
outputs, state = self.rnn(inputs, begin_state)
# assert_shape(outputs, (-1, unroll_length, self.num_cells))
# for s in state:
# assert_shape(s, (-1, self.num_cells))
# assert_shape(
# lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)),
# )
return outputs, state, lags_scaled, inputs
def unroll_encoder(
self,
past_time_feat: torch.Tensor,
past_target_cdf: torch.Tensor,
past_observed_values: torch.Tensor,
past_is_pad: torch.Tensor,
future_time_feat: Optional[torch.Tensor],
future_target_cdf: Optional[torch.Tensor],
target_dimension_indicator: torch.Tensor,
) -> Tuple[
torch.Tensor,
Union[List[torch.Tensor], torch.Tensor],
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""
Unrolls the RNN encoder over past and, if present, future data.
Returns outputs and state of the encoder, plus the scale of
past_target_cdf and a vector of static features that was constructed
and fed as input to the encoder. All tensor arguments should have NTC
layout.
Parameters
----------
past_time_feat
Past time features (batch_size, history_length, num_features)
past_target_cdf
Past marginal CDF transformed target values (batch_size,
history_length, target_dim)
past_observed_values
Indicator whether or not the values were observed (batch_size,
history_length, target_dim)
past_is_pad
Indicator whether the past target values have been padded
(batch_size, history_length)
future_time_feat
Future time features (batch_size, prediction_length, num_features)
future_target_cdf
Future marginal CDF transformed target values (batch_size,
prediction_length, target_dim)
target_dimension_indicator
Dimensionality of the time series (batch_size, target_dim)
Returns
-------
outputs
RNN outputs (batch_size, seq_len, num_cells)
states
RNN states. Nested list with (batch_size, num_cells) tensors with
dimensions target_dim x num_layers x (batch_size, num_cells)
scale
Mean scales for the time series (batch_size, 1, target_dim)
lags_scaled
Scaled lags(batch_size, sub_seq_len, target_dim, num_lags)
inputs
inputs to the RNN
"""
past_observed_values = torch.min(
past_observed_values, 1 - past_is_pad.unsqueeze(-1)
)
if future_time_feat is None or future_target_cdf is None:
time_feat = past_time_feat[:, -self.context_length :, ...]
sequence = past_target_cdf
sequence_length = self.history_length
subsequences_length = self.context_length
else:
time_feat = torch.cat(
(past_time_feat[:, -self.context_length :, ...], future_time_feat),
dim=1,
)
sequence = torch.cat((past_target_cdf, future_target_cdf), dim=1)
sequence_length = self.history_length + self.prediction_length
subsequences_length = self.context_length + self.prediction_length
# (batch_size, sub_seq_len, target_dim, num_lags)
lags = self.get_lagged_subsequences(
sequence=sequence,
sequence_length=sequence_length,
indices=self.lags_seq,
subsequences_length=subsequences_length,
)
# scale is computed on the context length last units of the past target
# scale shape is (batch_size, 1, target_dim)
_, scale = self.scaler(
past_target_cdf[:, -self.context_length :, ...],
past_observed_values[:, -self.context_length :, ...],
)
outputs, states, lags_scaled, inputs = self.unroll(
lags=lags,
scale=scale,
time_feat=time_feat,
target_dimension_indicator=target_dimension_indicator,
unroll_length=subsequences_length,
begin_state=None,
)
return outputs, states, scale, lags_scaled, inputs
def distr_args(self, rnn_outputs: torch.Tensor):
"""
Returns the distribution of DeepVAR with respect to the RNN outputs.
Parameters
----------
rnn_outputs
Outputs of the unrolled RNN (batch_size, seq_len, num_cells)
scale
Mean scale for each time series (batch_size, 1, target_dim)
Returns
-------
distr
Distribution instance
distr_args
Distribution arguments
"""
(distr_args,) = self.proj_dist_args(rnn_outputs)
# # compute likelihood of target given the predicted parameters
# distr = self.distr_output.distribution(distr_args, scale=scale)
# return distr, distr_args
return distr_args
def forward(
self,
target_dimension_indicator: torch.Tensor,
past_time_feat: torch.Tensor,
past_target_cdf: torch.Tensor,
past_observed_values: torch.Tensor,
past_is_pad: torch.Tensor,
future_time_feat: torch.Tensor,
future_target_cdf: torch.Tensor,
future_observed_values: torch.Tensor,
) -> Tuple[torch.Tensor, ...]:
"""
Computes the loss for training DeepVAR, all inputs tensors representing
time series have NTC layout.
Parameters
----------
target_dimension_indicator
Indices of the target dimension (batch_size, target_dim)
past_time_feat
Dynamic features of past time series (batch_size, history_length,
num_features)
past_target_cdf
Past marginal CDF transformed target values (batch_size,
history_length, target_dim)
past_observed_values
Indicator whether or not the values were observed (batch_size,
history_length, target_dim)
past_is_pad
Indicator whether the past target values have been padded
(batch_size, history_length)
future_time_feat
Future time features (batch_size, prediction_length, num_features)
future_target_cdf
Future marginal CDF transformed target values (batch_size,
prediction_length, target_dim)
future_observed_values
Indicator whether or not the future values were observed
(batch_size, prediction_length, target_dim)
Returns
-------
distr
Loss with shape (batch_size, 1)
likelihoods
Likelihoods for each time step
(batch_size, context + prediction_length, 1)
distr_args
Distribution arguments (context + prediction_length,
number_of_arguments)
"""
seq_len = self.context_length + self.prediction_length
# unroll the decoder in "training mode", i.e. by providing future data
# as well
rnn_outputs, _, scale, _, _ = self.unroll_encoder(
past_time_feat=past_time_feat,
past_target_cdf=past_target_cdf,
past_observed_values=past_observed_values,
past_is_pad=past_is_pad,
future_time_feat=future_time_feat,
future_target_cdf=future_target_cdf,
target_dimension_indicator=target_dimension_indicator,
)
# put together target sequence
# (batch_size, seq_len, target_dim)
target = torch.cat(
(past_target_cdf[:, -self.context_length :, ...], future_target_cdf),
dim=1,
)
# assert_shape(target, (-1, seq_len, self.target_dim))
distr_args = self.distr_args(rnn_outputs=rnn_outputs)
if self.scaling:
self.diffusion.scale = scale
# we sum the last axis to have the same shape for all likelihoods
# (batch_size, subseq_length, 1)
likelihoods = self.diffusion.log_prob(target, distr_args).unsqueeze(-1)
# assert_shape(likelihoods, (-1, seq_len, 1))
past_observed_values = torch.min(
past_observed_values, 1 - past_is_pad.unsqueeze(-1)
)
# (batch_size, subseq_length, target_dim)
observed_values = torch.cat(
(
past_observed_values[:, -self.context_length :, ...],
future_observed_values,
),
dim=1,
)
# mask the loss at one time step if one or more observations is missing
# in the target dimensions (batch_size, subseq_length, 1)
loss_weights, _ = observed_values.min(dim=-1, keepdim=True)
# assert_shape(loss_weights, (-1, seq_len, 1))
loss = weighted_average(likelihoods, weights=loss_weights, dim=1)
# assert_shape(loss, (-1, -1, 1))
# self.distribution = distr
return (loss.mean(), likelihoods, distr_args)
class TimeGradPredictionNetwork(TimeGradTrainingNetwork):
def __init__(self, num_parallel_samples: int, **kwargs) -> None:
super().__init__(**kwargs)
self.num_parallel_samples = num_parallel_samples
# for decoding the lags are shifted by one,
# at the first time-step of the decoder a lag of one corresponds to
# the last target value
self.shifted_lags = [l - 1 for l in self.lags_seq]
def sampling_decoder(
self,
past_target_cdf: torch.Tensor,
target_dimension_indicator: torch.Tensor,
time_feat: torch.Tensor,
scale: torch.Tensor,
begin_states: Union[List[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
"""
Computes sample paths by unrolling the RNN starting with a initial
input and state.
Parameters
----------
past_target_cdf
Past marginal CDF transformed target values (batch_size,
history_length, target_dim)
target_dimension_indicator
Indices of the target dimension (batch_size, target_dim)
time_feat
Dynamic features of future time series (batch_size, history_length,
num_features)
scale
Mean scale for each time series (batch_size, 1, target_dim)
begin_states
List of initial states for the RNN layers (batch_size, num_cells)
Returns
--------
sample_paths : Tensor
A tensor containing sampled paths. Shape: (1, num_sample_paths,
prediction_length, target_dim).
"""
def repeat(tensor, dim=0):
return tensor.repeat_interleave(repeats=self.num_parallel_samples, dim=dim)
# blows-up the dimension of each tensor to
# batch_size * self.num_sample_paths for increasing parallelism
repeated_past_target_cdf = repeat(past_target_cdf)
repeated_time_feat = repeat(time_feat)
repeated_scale = repeat(scale)
if self.scaling:
self.diffusion.scale = repeated_scale
repeated_target_dimension_indicator = repeat(target_dimension_indicator)
if self.cell_type == "LSTM":
repeated_states = [repeat(s, dim=1) for s in begin_states]
else:
repeated_states = repeat(begin_states, dim=1)
future_samples = []
# for each future time-units we draw new samples for this time-unit
# and update the state
for k in range(self.prediction_length):
lags = self.get_lagged_subsequences(
sequence=repeated_past_target_cdf,
sequence_length=self.history_length + k,
indices=self.shifted_lags,
subsequences_length=1,
)
rnn_outputs, repeated_states, _, _ = self.unroll(
begin_state=repeated_states,
lags=lags,
scale=repeated_scale,
time_feat=repeated_time_feat[:, k : k + 1, ...],
target_dimension_indicator=repeated_target_dimension_indicator,
unroll_length=1,
)
distr_args = self.distr_args(rnn_outputs=rnn_outputs)
# (batch_size, 1, target_dim)
new_samples = self.diffusion.sample(cond=distr_args)
# (batch_size, seq_len, target_dim)
future_samples.append(new_samples)
repeated_past_target_cdf = torch.cat(
(repeated_past_target_cdf, new_samples), dim=1
)
# (batch_size * num_samples, prediction_length, target_dim)
samples = torch.cat(future_samples, dim=1)
# (batch_size, num_samples, prediction_length, target_dim)
return samples.reshape(
(
-1,
self.num_parallel_samples,
self.prediction_length,
self.target_dim,
)
)
def forward(
self,
target_dimension_indicator: torch.Tensor,
past_time_feat: torch.Tensor,
past_target_cdf: torch.Tensor,
past_observed_values: torch.Tensor,
past_is_pad: torch.Tensor,
future_time_feat: torch.Tensor,
) -> torch.Tensor:
"""
Predicts samples given the trained DeepVAR model.
All tensors should have NTC layout.
Parameters
----------
target_dimension_indicator
Indices of the target dimension (batch_size, target_dim)
past_time_feat
Dynamic features of past time series (batch_size, history_length,
num_features)
past_target_cdf
Past marginal CDF transformed target values (batch_size,
history_length, target_dim)
past_observed_values
Indicator whether or not the values were observed (batch_size,
history_length, target_dim)
past_is_pad
Indicator whether the past target values have been padded
(batch_size, history_length)
future_time_feat
Future time features (batch_size, prediction_length, num_features)
Returns
-------
sample_paths : Tensor
A tensor containing sampled paths (1, num_sample_paths,
prediction_length, target_dim).
"""
# mark padded data as unobserved
# (batch_size, target_dim, seq_len)
past_observed_values = torch.min(
past_observed_values, 1 - past_is_pad.unsqueeze(-1)
)
# unroll the decoder in "prediction mode", i.e. with past data only
_, begin_states, scale, _, _ = self.unroll_encoder(
past_time_feat=past_time_feat,
past_target_cdf=past_target_cdf,
past_observed_values=past_observed_values,
past_is_pad=past_is_pad,
future_time_feat=None,
future_target_cdf=None,
target_dimension_indicator=target_dimension_indicator,
)
return self.sampling_decoder(
past_target_cdf=past_target_cdf,
target_dimension_indicator=target_dimension_indicator,
time_feat=future_time_feat,
scale=scale,
begin_states=begin_states,
)
pts/modules/__init__.py
+2
View File
@@ -13,8 +13,10 @@ from .distribution_output import (
LowRankMultivariateNormalOutput,
MultivariateNormalOutput,
FlowOutput,
DiffusionOutput,
ImplicitQuantileOutput,
)
from .feature import FeatureEmbedder, FeatureAssembler
from .flows import RealNVP, MAF
from .scaler import MeanScaler, NOPScaler
from .gaussian_diffusion import GaussianDiffusion
pts/modules/distribution_output.py
+24
View File
@@ -449,6 +449,30 @@ class FlowOutput(DistributionOutput):
return (self.dim,)
class DiffusionOutput(DistributionOutput):
@validated()
def __init__(self, diffusion, input_size, cond_size):
self.args_dim = {"cond": cond_size}
self.diffusion = diffusion
self.dim = input_size
@classmethod
def domain_map(cls, cond):
return (cond,)
def distribution(self, distr_args, scale=None):
(cond,) = distr_args
if scale is not None:
self.diffusion.scale = scale
self.diffusion.cond = cond
return self.diffusion
@property
def event_shape(self) -> Tuple:
return (self.dim,)
class QuantilePtArgProj(PtArgProj):
def __init__(
self,
pts/modules/gaussian_diffusion.py
+269
View File
@@ -0,0 +1,269 @@
from functools import partial
from inspect import isfunction
import numpy as np
import torch
from torch import nn, einsum
import torch.nn.functional as F
def default(val, d):
if val is not None:
return val
return d() if isfunction(d) else d
def extract(a, t, x_shape):
b, *_ = t.shape
out = a.gather(-1, t)
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
def noise_like(shape, device, repeat=False):
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(
shape[0], *((1,) * (len(shape) - 1))
)
noise = lambda: torch.randn(shape, device=device)
return repeat_noise() if repeat else noise()
def cosine_beta_schedule(timesteps, s=0.008):
"""
cosine schedule
as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
"""
steps = timesteps + 1
x = np.linspace(0, steps, steps)
alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
return np.clip(betas, a_min=0, a_max=0.999)
class GaussianDiffusion(nn.Module):
def __init__(
self,
denoise_fn,
input_size,
beta_end=0.1,
diff_steps=100,
loss_type="l2",
betas=None,
beta_schedule="linear",
):
super().__init__()
self.denoise_fn = denoise_fn
self.input_size = input_size
self.__scale = None
if betas is not None:
betas = (
betas.detach().cpu().numpy()
if isinstance(betas, torch.Tensor)
else betas
)
else:
if beta_schedule == "linear":
betas = np.linspace(1e-4, beta_end, diff_steps)
elif beta_schedule == "quad":
betas = np.linspace(1e-4 ** 0.5, beta_end ** 0.5, diff_steps) ** 2
elif beta_schedule == "const":
betas = beta_end * np.ones(diff_steps)
elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1
betas = 1.0 / np.linspace(diff_steps, 1, diff_steps)
elif beta_schedule == "sigmoid":
betas = np.linspace(-6, 6, diff_steps)
betas = (beta_end - 1e-4) / (np.exp(-betas) + 1) + 1e-4
elif beta_schedule == "cosine":
betas = cosine_beta_schedule(diff_steps)
else:
raise NotImplementedError(beta_schedule)
alphas = 1.0 - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])
(timesteps,) = betas.shape
self.num_timesteps = int(timesteps)
self.loss_type = loss_type
to_torch = partial(torch.tensor, dtype=torch.float32)
self.register_buffer("betas", to_torch(betas))
self.register_buffer("alphas_cumprod", to_torch(alphas_cumprod))
self.register_buffer("alphas_cumprod_prev", to_torch(alphas_cumprod_prev))
# calculations for diffusion q(x_t | x_{t-1}) and others
self.register_buffer("sqrt_alphas_cumprod", to_torch(np.sqrt(alphas_cumprod)))
self.register_buffer(
"sqrt_one_minus_alphas_cumprod", to_torch(np.sqrt(1.0 - alphas_cumprod))
)
self.register_buffer(
"log_one_minus_alphas_cumprod", to_torch(np.log(1.0 - alphas_cumprod))
)
self.register_buffer(
"sqrt_recip_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod))
)
self.register_buffer(
"sqrt_recipm1_alphas_cumprod", to_torch(np.sqrt(1.0 / alphas_cumprod - 1))
)
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_variance = (
betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
)
# above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
self.register_buffer("posterior_variance", to_torch(posterior_variance))
# below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
self.register_buffer(
"posterior_log_variance_clipped",
to_torch(np.log(np.maximum(posterior_variance, 1e-20))),
)
self.register_buffer(
"posterior_mean_coef1",
to_torch(betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)),
)
self.register_buffer(
"posterior_mean_coef2",
to_torch(
(1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod)
),
)
@property
def scale(self):
return self.__scale
@scale.setter
def scale(self, scale):
self.__scale = scale
def q_mean_variance(self, x_start, t):
mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
variance = extract(1.0 - self.alphas_cumprod, t, x_start.shape)
log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)
return mean, variance, log_variance
def predict_start_from_noise(self, x_t, t, noise):
return (
extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
- extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
)
def q_posterior(self, x_start, x_t, t):
posterior_mean = (
extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
+ extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
)
posterior_variance = extract(self.posterior_variance, t, x_t.shape)
posterior_log_variance_clipped = extract(
self.posterior_log_variance_clipped, t, x_t.shape
)
return posterior_mean, posterior_variance, posterior_log_variance_clipped
def p_mean_variance(self, x, cond, t, clip_denoised: bool):
x_recon = self.predict_start_from_noise(
x, t=t, noise=self.denoise_fn(x, t, cond=cond)
)
if clip_denoised:
x_recon.clamp_(-1.0, 1.0)
model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
x_start=x_recon, x_t=x, t=t
)
return model_mean, posterior_variance, posterior_log_variance
@torch.no_grad()
def p_sample(self, x, cond, t, clip_denoised=False, repeat_noise=False):
b, *_, device = *x.shape, x.device
model_mean, _, model_log_variance = self.p_mean_variance(
x=x, cond=cond, t=t, clip_denoised=clip_denoised
)
noise = noise_like(x.shape, device, repeat_noise)
# no noise when t == 0
nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
@torch.no_grad()
def p_sample_loop(self, shape, cond):
device = self.betas.device
b = shape[0]
img = torch.randn(shape, device=device)
for i in reversed(range(0, self.num_timesteps)):
img = self.p_sample(
img, cond, torch.full((b,), i, device=device, dtype=torch.long)
)
return img
@torch.no_grad()
def sample(self, sample_shape=torch.Size(), cond=None):
if cond is not None:
shape = cond.shape[:-1] + (self.input_size,)
# TODO reshape cond to (B*T, 1, -1)
else:
shape = sample_shape
x_hat = self.p_sample_loop(shape, cond) # TODO reshape x_hat to (B,T,-1)
if self.scale is not None:
x_hat *= self.scale
return x_hat
@torch.no_grad()
def interpolate(self, x1, x2, t=None, lam=0.5):
b, *_, device = *x1.shape, x1.device
t = default(t, self.num_timesteps - 1)
assert x1.shape == x2.shape
t_batched = torch.stack([torch.tensor(t, device=device)] * b)
xt1, xt2 = map(lambda x: self.q_sample(x, t=t_batched), (x1, x2))
img = (1 - lam) * xt1 + lam * xt2
for i in reversed(range(0, t)):
img = self.p_sample(
img, torch.full((b,), i, device=device, dtype=torch.long)
)
return img
def q_sample(self, x_start, t, noise=None):
noise = default(noise, lambda: torch.randn_like(x_start))
return (
extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+ extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
)
def p_losses(self, x_start, cond, t, noise=None):
noise = default(noise, lambda: torch.randn_like(x_start))
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
x_recon = self.denoise_fn(x_noisy, t, cond=cond)
if self.loss_type == "l1":
loss = F.l1_loss(x_recon, noise)
elif self.loss_type == "l2":
loss = F.mse_loss(x_recon, noise)
elif self.loss_type == "huber":
loss = F.smooth_l1_loss(x_recon, noise)
else:
raise NotImplementedError()
return loss
def log_prob(self, x, cond, *args, **kwargs):
if self.scale is not None:
x /= self.scale
B, T, _ = x.shape
time = torch.randint(0, self.num_timesteps, (B * T,), device=x.device).long()
loss = self.p_losses(
x.reshape(B * T, 1, -1), cond.reshape(B * T, 1, -1), time, *args, **kwargs
)
return loss
pts/trainer.py
+2 -5
View File
@@ -1,6 +1,5 @@
import time
from typing import List, Optional, Union
from torch.optim import lr_scheduler
from tqdm import tqdm
import wandb
@@ -47,9 +46,7 @@ class Trainer:
wandb.watch(net, log="all", log_freq=self.num_batches_per_epoch)
optimizer = Adam(
net.parameters(),
lr=self.learning_rate,
weight_decay=self.weight_decay
net.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay
)
lr_scheduler = OneCycleLR(
@@ -88,7 +85,7 @@ class Trainer:
loss.backward()
if self.clip_gradient is not None:
nn.utils.clip_grad_norm_(net.parameters(), self.clip_gradient)
optimizer.step()
lr_scheduler.step()
setup.py
+2 -3
View File
@@ -18,11 +18,10 @@ setup(
"torch>=1.7.0",
"gluonts",
"holidays",
"numpy",
"pandas>=1.1",
"numpy~=1.16",
"pandas~=1.1",
"scipy",
"tqdm",
"pydantic",
"matplotlib",
"tensorboard",
"wandb",