File size: 4,239 Bytes
bdd9175 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """TinyCast model assembly.
model(past_values, scale_factor, prediction_length, batch_first).quantile_outputs
"""
from dataclasses import dataclass
import torch
import torch.nn as nn
from .config import TinyCastConfig
@dataclass
class PredictionOutput:
quantile_outputs: torch.Tensor
prediction_outputs: torch.Tensor = None
class TinyCastBackbone(nn.Module):
"""Inner model: per-window min-max norm -> dilated-conv core -> denorm."""
def __init__(self, config: TinyCastConfig):
super().__init__()
self.config = config
from .backbone import DilatedConvBackbone
from .normalization import WindowMinMax
self.core = DilatedConvBackbone(
seq_len=int(config.seq_len),
p_out=int(config.output_token_len),
n_quantiles=int(config.num_quantiles),
d=int(config.conv_dim),
n_layers=int(config.n_layers),
kernel=int(config.kernel_size),
ffn_mult=float(config.ffn_mult),
top_k_periods=int(config.top_k_periods),
significance_alpha=float(config.significance_alpha),
n_harmonics=int(config.n_harmonics),
pool_kind=str(config.pool_kind),
causal=bool(config.causal),
phase_bins=int(config.phase_bins),
decoder_depth=int(config.decoder_depth),
separable_conv=bool(config.separable_conv),
share_ffn=bool(config.share_ffn),
future_conv=bool(config.future_conv),
future_conv_layers=int(config.future_conv_layers),
future_conv_seed=int(config.future_conv_seed),
)
self.norm = WindowMinMax(eps_clamp=1e-5)
def encode(
self,
past_values: torch.Tensor,
batch_first: bool = True,
scale_factor: "torch.Tensor | float | None" = None,
horizon: "int | None" = None,
) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]":
"""per-window min-max norm -> dilated-conv core -> normalized y."""
x = past_values
if not batch_first:
x = x.transpose(0, 1)
if x.dim() == 2:
x = x.unsqueeze(-1)
x_normed, x_min, x_range = self.norm.transform(x)
nan_mask = (~torch.isnan(x)).to(x.dtype)
y_norm = self.core(
x_normed, nan_mask=nan_mask, scale_factor=scale_factor,
horizon=horizon,
)
return y_norm, x_min, x_range
class TinyCastForPrediction(nn.Module):
"""Top-level model container.
Usage:
model = TinyCastForPrediction(TinyCastConfig())
out = model(past_values=x, scale_factor=sf, prediction_length=pl,
batch_first=False)
quantiles = out.quantile_outputs # (B, Q, pred_len, 1)
NOTE: this ``forward`` is the SINGLE-SHOT arbitrary-horizon path. The
deployed GIFT-Eval inference uses the AR-rollout predictor
(``tinycast.predictor.TinyCastPredictor``), which calls ``self.model.encode``
in 48-step chunks. For horizons <= 48 the two paths coincide.
"""
def __init__(self, config: TinyCastConfig):
super().__init__()
self.config = config
self.model = TinyCastBackbone(config)
def forward(self, past_values, scale_factor=None, prediction_length=None,
batch_first=None):
if scale_factor is None:
scale_factor = 1.0
if batch_first is None:
batch_first = True
ctx = past_values
if not batch_first:
ctx = ctx.transpose(0, 1)
if ctx.dim() == 2:
ctx = ctx.unsqueeze(-1)
ctx_len_max = int(self.config.seq_len)
ctx_in = ctx[:, -ctx_len_max:, :]
p_out_native = int(self.model.core.p_out)
pred_len = int(prediction_length) if prediction_length else p_out_native
y_norm, x_min, x_range = self.model.encode(
ctx_in, batch_first=True, scale_factor=scale_factor,
horizon=pred_len,
)
if y_norm.dim() == 2:
y_norm = y_norm.unsqueeze(-1)
y_pred = y_norm * x_range + x_min
quantile_outputs = y_pred.permute(0, 2, 1).unsqueeze(-1)
return PredictionOutput(quantile_outputs=quantile_outputs)
|