File size: 3,055 Bytes
bffb943 | 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 | """
A small, patch-based transformer for zero-shot time series forecasting —
a toy-scale version of the same general recipe used by real time-series
foundation models (Chronos, TimesFM, t0-alpha, etc.):
1. Per-series normalization: each input window is normalized by its
OWN mean and std before entering the model, and de-normalized
using those same stats afterward. This is the single most
important design choice for generalization — it means the model
learns SHAPE (trend, seasonality, noise structure), not absolute
numbers, so a model trained on synthetic series scaled 0.5-500 can
meaningfully forecast a real series scaled in the thousands.
2. Patching: rather than feeding every single timestep as a separate
token (expensive, and treats each point independently), the input
window is split into small chunks ("patches") and each patch is
projected into the model's embedding space as one token. This is
the same idea PatchTST and most modern time-series transformers
use — it's both more efficient and tends to generalize better
than point-wise tokenization.
3. A standard transformer encoder processes the patch tokens, and a
linear head maps the final representation directly to the
forecast horizon.
CONTEXT_LEN and HORIZON are fixed (96 and 42) to match the evaluation
setup already used for SARIMA/Prophet/statsforecast/t0-alpha in the
sibling demand-forecasting project, so results are directly comparable.
"""
import torch
import torch.nn as nn
CONTEXT_LEN = 96
HORIZON = 42
PATCH_SIZE = 8
N_PATCHES = CONTEXT_LEN // PATCH_SIZE # 12
D_MODEL = 32
N_HEADS = 4
N_LAYERS = 2
DIM_FEEDFORWARD = 64
class TinyForecaster(nn.Module):
def __init__(self):
super().__init__()
self.patch_embed = nn.Linear(PATCH_SIZE, D_MODEL)
self.pos_embed = nn.Parameter(torch.randn(1, N_PATCHES, D_MODEL) * 0.02)
encoder_layer = nn.TransformerEncoderLayer(
d_model=D_MODEL, nhead=N_HEADS, dim_feedforward=DIM_FEEDFORWARD,
batch_first=True, dropout=0.1,
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=N_LAYERS)
self.head = nn.Sequential(
nn.Linear(D_MODEL * N_PATCHES, 128),
nn.ReLU(),
nn.Linear(128, HORIZON),
)
def forward(self, context_raw: torch.Tensor) -> torch.Tensor:
"""
context_raw: (batch, CONTEXT_LEN) raw (unnormalized) values.
Returns: (batch, HORIZON) forecast, in the ORIGINAL scale.
"""
mean = context_raw.mean(dim=1, keepdim=True)
std = context_raw.std(dim=1, keepdim=True).clamp(min=1e-3)
normed = (context_raw - mean) / std
patches = normed.view(-1, N_PATCHES, PATCH_SIZE)
tokens = self.patch_embed(patches) + self.pos_embed
encoded = self.encoder(tokens)
flat = encoded.reshape(encoded.size(0), -1)
forecast_normed = self.head(flat)
forecast = forecast_normed * std + mean
return forecast
|