LOTUS / src /opentslm /model /decoder /TransformerCNNDecoder2D.py

Add load_lotus.py one-import loader with vendored model code

13a7cc8 verified 2 days ago

4.52 kB

# 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."
        )

	# 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."
	)