File size: 4,518 Bytes
13a7cc8 | 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 120 121 122 123 124 | # SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT
import torch
import torch.nn as nn
from opentslm.model_config import PATCH_SIZE
from opentslm.model.decoder.TimeSeriesDecoderBase import TimeSeriesDecoderBase
class TransformerCNNDecoder2D(TimeSeriesDecoderBase):
"""
Causal numeric decoder over encoder-projected patch states.
This module assumes a single forecasting workflow:
history -> encoder -> projector -> decoder -> next patch
Input shapes accepted:
(N, T, input_dim) → (N, T, patch_size) full sequence
(N, input_dim) → (N, patch_size) single-step convenience
"""
def __init__(
self,
input_dim: int = 768, # enc_dim — output of decoder_projector
decoder_dim: int = 512, # internal width of the causal transformer - OLD: 512
patch_size: int = PATCH_SIZE,
dropout: float = 0.1,
num_heads: int = 8,
num_layers: int = 4, # OLD: 4
ff_dim: int = 2048, # OLD: 2048
# legacy / unused kwargs kept for API compatibility
pe_base: float = 10000.0,
use_conv_decoder: bool = True,
**kwargs,
):
super().__init__(input_dim, patch_size, dropout)
self.input_dim = input_dim
self.decoder_dim = decoder_dim
# Project from enc_dim → decoder_dim (Identity if equal)
self.in_proj = (
nn.Linear(input_dim, decoder_dim)
if input_dim != decoder_dim
else nn.Identity()
)
# GPT-style causal self-attention stack (no cross-attention)
encoder_layer = nn.TransformerEncoderLayer(
d_model=decoder_dim,
nhead=num_heads,
dim_feedforward=ff_dim,
dropout=dropout,
batch_first=True,
activation="gelu",
norm_first=True, # pre-LN for stable deep training
)
self.transformer = nn.TransformerEncoder(
encoder_layer,
num_layers=num_layers,
norm=nn.LayerNorm(decoder_dim),
)
# Patch prediction head
self.head = nn.Sequential(
nn.LayerNorm(decoder_dim),
nn.Linear(decoder_dim, patch_size),
)
self._print_param_count()
def _print_param_count(self):
total = sum(p.numel() for p in self.parameters())
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
print(f"🔢 Decoder: {total:,} params ({trainable:,} trainable)")
def _decode_hidden(self, x: torch.Tensor):
"""Run the causal transformer stack and return hidden states."""
squeeze = x.dim() == 2
if squeeze:
x = x.unsqueeze(1) # (N, 1, input_dim)
_, T, _ = x.shape
out = self.in_proj(x) # (N, T, decoder_dim)
causal_mask = nn.Transformer.generate_square_subsequent_mask(
T, device=x.device, dtype=out.dtype
)
out = self.transformer(out, mask=causal_mask, is_causal=True) # (N, T, decoder_dim)
return out, squeeze
def forward(self, x: torch.Tensor, chan_id=None) -> torch.Tensor:
"""
x : (N, T, input_dim) — sequence of projected LLM hidden states
(N, input_dim) — single hidden vector (inference convenience)
chan_id: unused, kept for API compatibility.
Returns:
(N, T, patch_size) when x is 3-D
(N, patch_size) when x is 2-D
"""
out, squeeze = self._decode_hidden(x)
out = self.head(out) # (N, T, patch_size)
if squeeze:
return out[:, 0, :] # (N, patch_size)
return out # (N, T, patch_size)
def forecast_teacher_forced(self, *args, **kwargs) -> torch.Tensor:
raise NotImplementedError(
"Teacher-forced decoder rollout has been removed. "
"Use the encoder -> projector -> decoder next-patch workflow instead."
)
def forecast_autoregressive(self, *args, **kwargs) -> torch.Tensor:
raise NotImplementedError(
"Decoder-internal autoregressive rollout has been removed. "
"Use an outer loop that re-encodes predicted patches through the model workflow."
)
|