LOTUS / src /opentslm /model /encoder /MLPEncoder2D.py

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

13a7cc8 verified 2 days ago

11.6 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

"""Lightweight MLP-based 2D patch encoder.
Drop-in replacement for ``TransformerCNNEncoder2D`` with the same
``forward(x, chan_id, attention_mask=None) -> (C, T, output_dim)`` signature.
Patch embedding and 2D sin-cos positional encoding are identical to the
transformer variant, but the contextual stack is replaced with a position-wise
2-layer MLP (no attention, no cross-patch mixing).
"""

from typing import List
import torch
import torch.nn as nn

from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
from opentslm.model.encoder.TimeSeriesEncoderBase import TimeSeriesEncoderBase


def _get_1d_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, base: float = 10000.0) -> torch.Tensor:
    """MAE-style 1D sin-cos embedding. pos: (M,) float → (M, embed_dim)."""
    assert embed_dim % 2 == 0
    device, dtype = pos.device, pos.dtype
    omega = torch.arange(embed_dim // 2, device=device, dtype=dtype)
    omega = omega / (embed_dim / 2.0)
    omega = 1.0 / (base ** omega)
    out = pos.unsqueeze(1) * omega.unsqueeze(0)
    return torch.cat([torch.sin(out), torch.cos(out)], dim=1)


class MLPEncoder2D(TimeSeriesEncoderBase):
    """Position-wise MLP encoder with 2D (time, channel) sin-cos PE.
    Input : x of shape (C, L), chan_id of shape (C,), optional attention_mask (C, T).
    Output: (C, T, output_dim), same shape contract as TransformerCNNEncoder2D.
    The contextual stack is a 2-layer MLP applied independently to each patch
    embedding. Useful as a lightweight baseline.
    """

    def __init__(
        self,
        args=None,
        output_dim: int = ENCODER_OUTPUT_DIM,
        dropout: float = 0.1,
        transformer_input_dim: int = TRANSFORMER_INPUT_DIM,
        patch_size: int = PATCH_SIZE,
        pe_base: float = 10000.0,        # legacy single-base fallback
        pe_base_t: float | None = None,  # overrides pe_base for TIME axis
        pe_base_c: float | None = None,  # overrides pe_base for CHANNEL axis
        ff_dim: int = 3072,
    ):
        super().__init__(output_dim, dropout)
        self.patch_size = patch_size
        # Separate bases for time vs channel. If `args` carries them (via CLI),
        # prefer those; otherwise fall back to constructor kwargs; finally to
        # `pe_base` for backward compatibility.
        _pe_t_arg = getattr(args, "pe_base_t", None) if args is not None else None
        _pe_c_arg = getattr(args, "pe_base_c", None) if args is not None else None
        self.pe_base_t = float(_pe_t_arg if _pe_t_arg is not None else (pe_base_t if pe_base_t is not None else pe_base))
        self.pe_base_c = float(_pe_c_arg if _pe_c_arg is not None else (pe_base_c if pe_base_c is not None else pe_base))
        # `pe_base` kept for legacy compatibility (callers that still reference
        # `self.pe_base`).
        self.pe_base = float(pe_base)
        self.transformer_input_dim = transformer_input_dim
        self.args = args

        # Patch embedding: (C, 1, L) → (C, D, T) via strided conv over time.
        self.patch_embed = nn.Conv1d(
            in_channels=1,
            out_channels=transformer_input_dim,
            kernel_size=patch_size,
            stride=patch_size,
            bias=False,
        )

        # Frequency patch embedding (optional). Mirrors TransformerCNNEncoder2D:
        # rfft on a patch of length P gives P//2+1 complex values → 2*(P//2+1) reals,
        # which are projected to transformer_input_dim then concat-and-projected with
        # the time-domain patch embedding.
        self._use_freq_embed = bool(getattr(args, "use_freq_embed", False))
        # How the frequency branch encodes each patch spectrum.
        #   "linear" (default, legacy): concat(real, imag) -> bias-free Linear.
        #       NOTE: this makes the whole branch an exact linear reparameterisation
        #       of `patch_embed`. rfft is a linear map, and patch_embed / freq_embed /
        #       concat_proj are all bias-free with no nonlinearity between them, so
        #       the fused output equals a single Linear(P -> D) of the raw patch.
        #       It adds parameters, not expressivity. Kept as the default so that
        #       existing `*_freq` checkpoints keep computing what they were trained to.
        #   "spectral": concat(magnitude, phase) -> Linear -> GELU. |F(x)| is
        #       nonlinear in x, so this genuinely enlarges the function class.
        #       Shapes and parameter counts are identical, so a legacy checkpoint
        #       will LOAD SILENTLY but be MISINTERPRETED -- only use this mode with
        #       a model trained under it.
        self._freq_mode = str(getattr(args, "freq_mode", "linear")).lower()
        if self._freq_mode not in ("linear", "spectral"):
            raise ValueError(
                f"freq_mode must be 'linear' or 'spectral', got {self._freq_mode!r}"
            )
        if self._use_freq_embed:
            freq_input_dim = (patch_size // 2 + 1) * 2
            self.freq_embed = nn.Linear(freq_input_dim, transformer_input_dim, bias=False)
            self.concat_proj = nn.Linear(2 * transformer_input_dim, transformer_input_dim, bias=False)

        self.input_norm = nn.LayerNorm(transformer_input_dim)
        self.input_dropout = nn.Dropout(self.dropout)

        # Position-wise 2-layer MLP (same depth as requested).
        self.mlp = nn.Sequential(
            nn.Linear(transformer_input_dim, ff_dim),
            nn.GELU(),
            nn.Dropout(self.dropout),
            nn.Linear(ff_dim, transformer_input_dim),
        )
        self.out_norm = nn.LayerNorm(transformer_input_dim)

        if transformer_input_dim != output_dim:
            self.output_proj = nn.Linear(transformer_input_dim, output_dim)
        else:
            self.output_proj = nn.Identity()

        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"🔢 MLPEncoder2D: {total:,} params ({trainable:,} trainable)")

    # ------------------------------------------------------------------
    # Shared helpers
    # ------------------------------------------------------------------

    def _embed_and_pe(self, tok: torch.Tensor, chan_id: torch.Tensor, t_offset: int = 0) -> torch.Tensor:
        """Add 2D sin-cos PE to a (C, T, D) patch-embed tensor."""
        C, T, D = tok.shape
        t_pos = torch.arange(t_offset, t_offset + T, device=tok.device, dtype=torch.float32)
        c_pos = chan_id.to(device=tok.device, dtype=torch.float32)

        use_only_time_pe = bool(getattr(self.args, "use_only_time_pe", False))
        use_only_channel_pe = bool(getattr(self.args, "use_only_channel_pe", False))

        if use_only_time_pe:
            pe = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
            pe = pe.unsqueeze(0).expand(C, -1, -1)
        elif use_only_channel_pe:
            pe = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
            pe = pe.unsqueeze(1).expand(-1, T, -1)
        else:
            if D % 4 == 0:
                pe_t = _get_1d_sincos_pos_embed(D // 2, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D // 2, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = torch.cat(
                    [pe_t.unsqueeze(0).expand(C, -1, -1),
                     pe_c.unsqueeze(1).expand(-1, T, -1)],
                    dim=-1,
                )
            else:
                pe_t = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
                pe_c = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
                pe = pe_t.unsqueeze(0).expand(C, -1, -1) + pe_c.unsqueeze(1).expand(-1, T, -1)

        return tok + pe

    # ------------------------------------------------------------------
    # Main forward
    # ------------------------------------------------------------------

    def forward(
        self,
        x: torch.Tensor,
        chan_id: torch.Tensor,
        attention_mask: torch.Tensor = None,
    ) -> torch.Tensor:
        """x: (C, L), chan_id: (C,), attention_mask: (C, T) or None. Returns (C, T, output_dim)."""
        C, L = x.shape
        if L % self.patch_size != 0:
            raise ValueError(f"L={L} not divisible by patch_size={self.patch_size}")

        T = L // self.patch_size
        tok = self.patch_embed(x.reshape(C, 1, L)).transpose(1, 2)   # (C, T, D)

        if self._use_freq_embed:
            tok = self._fuse_freq(tok, x.reshape(C, T, self.patch_size))

        tok = self._embed_and_pe(tok, chan_id, t_offset=0)
        tok = self.input_norm(tok)
        tok = self.input_dropout(tok)

        # Position-wise MLP with residual.
        tok = tok + self.mlp(tok)
        tok = self.out_norm(tok)

        # Mask-out padding patches (zero their features) so downstream averaging
        # or projection doesn't leak padding content.
        if attention_mask is not None:
            m = attention_mask.to(dtype=tok.dtype, device=tok.device).unsqueeze(-1)
            tok = tok * m

        return self.output_proj(tok)   # (C, T, output_dim)

    # ------------------------------------------------------------------
    # Compatibility stubs for callers that expect the transformer's
    # incremental API. MLP is position-wise so "incremental" is just a
    # single patch forward; we reuse `forward` for simplicity.
    # ------------------------------------------------------------------

    def tokenize_patches(
        self,
        x: torch.Tensor,
        chan_id: torch.Tensor,
        t_offset: int = 0,
    ) -> torch.Tensor:
        if x.size(1) % self.patch_size != 0:
            raise ValueError(f"L={x.size(1)} not divisible by patch_size={self.patch_size}")
        C = x.size(0)
        T_new = x.size(1) // self.patch_size
        tok = self.patch_embed(x.reshape(C, 1, -1)).transpose(1, 2)

        if self._use_freq_embed:
            tok = self._fuse_freq(tok, x.reshape(C, T_new, self.patch_size))

        tok = self._embed_and_pe(tok, chan_id, t_offset=t_offset)
        return self.input_norm(tok)

    def _fuse_freq(self, tok: torch.Tensor, x_patches: torch.Tensor) -> torch.Tensor:
        """Fuse a per-patch spectral embedding into the time-domain patch tokens.
        tok:       (C, T, D) time-domain patch embeddings
        x_patches: (C, T, P) raw (z-normalised) patches
        Returns:   (C, T, D)
        See the `_freq_mode` note in __init__ for why "linear" is expressively
        inert and "spectral" is not.
        """
        x_fft = torch.fft.rfft(x_patches, dim=-1)              # (C, T, P//2+1) complex
        if self._freq_mode == "spectral":
            # Magnitude is nonlinear in x; phase supplies the shift information
            # that magnitude alone discards.
            feats = torch.cat([x_fft.abs(), x_fft.angle()], dim=-1)
            freq_tok = torch.nn.functional.gelu(self.freq_embed(feats.to(tok.dtype)))
        else:
            feats = torch.cat([x_fft.real, x_fft.imag], dim=-1)
            freq_tok = self.freq_embed(feats.to(tok.dtype))
        return self.concat_proj(torch.cat([tok, freq_tok], dim=-1))

	# 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

	"""Lightweight MLP-based 2D patch encoder.

	Drop-in replacement for ``TransformerCNNEncoder2D`` with the same
	``forward(x, chan_id, attention_mask=None) -> (C, T, output_dim)`` signature.
	Patch embedding and 2D sin-cos positional encoding are identical to the
	transformer variant, but the contextual stack is replaced with a position-wise
	2-layer MLP (no attention, no cross-patch mixing).
	"""

	from typing import List
	import torch
	import torch.nn as nn

	from opentslm.model_config import TRANSFORMER_INPUT_DIM, ENCODER_OUTPUT_DIM, PATCH_SIZE
	from opentslm.model.encoder.TimeSeriesEncoderBase import TimeSeriesEncoderBase


	def _get_1d_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, base: float = 10000.0) -> torch.Tensor:
	"""MAE-style 1D sin-cos embedding. pos: (M,) float → (M, embed_dim)."""
	assert embed_dim % 2 == 0
	device, dtype = pos.device, pos.dtype
	omega = torch.arange(embed_dim // 2, device=device, dtype=dtype)
	omega = omega / (embed_dim / 2.0)
	omega = 1.0 / (base ** omega)
	out = pos.unsqueeze(1) * omega.unsqueeze(0)
	return torch.cat([torch.sin(out), torch.cos(out)], dim=1)


	class MLPEncoder2D(TimeSeriesEncoderBase):
	"""Position-wise MLP encoder with 2D (time, channel) sin-cos PE.

	Input : x of shape (C, L), chan_id of shape (C,), optional attention_mask (C, T).
	Output: (C, T, output_dim), same shape contract as TransformerCNNEncoder2D.

	The contextual stack is a 2-layer MLP applied independently to each patch
	embedding. Useful as a lightweight baseline.
	"""

	def __init__(
	self,
	args=None,
	output_dim: int = ENCODER_OUTPUT_DIM,
	dropout: float = 0.1,
	transformer_input_dim: int = TRANSFORMER_INPUT_DIM,
	patch_size: int = PATCH_SIZE,
	pe_base: float = 10000.0, # legacy single-base fallback
	pe_base_t: float \| None = None, # overrides pe_base for TIME axis
	pe_base_c: float \| None = None, # overrides pe_base for CHANNEL axis
	ff_dim: int = 3072,
	):
	super().__init__(output_dim, dropout)
	self.patch_size = patch_size
	# Separate bases for time vs channel. If `args` carries them (via CLI),
	# prefer those; otherwise fall back to constructor kwargs; finally to
	# `pe_base` for backward compatibility.
	_pe_t_arg = getattr(args, "pe_base_t", None) if args is not None else None
	_pe_c_arg = getattr(args, "pe_base_c", None) if args is not None else None
	self.pe_base_t = float(_pe_t_arg if _pe_t_arg is not None else (pe_base_t if pe_base_t is not None else pe_base))
	self.pe_base_c = float(_pe_c_arg if _pe_c_arg is not None else (pe_base_c if pe_base_c is not None else pe_base))
	# `pe_base` kept for legacy compatibility (callers that still reference
	# `self.pe_base`).
	self.pe_base = float(pe_base)
	self.transformer_input_dim = transformer_input_dim
	self.args = args

	# Patch embedding: (C, 1, L) → (C, D, T) via strided conv over time.
	self.patch_embed = nn.Conv1d(
	in_channels=1,
	out_channels=transformer_input_dim,
	kernel_size=patch_size,
	stride=patch_size,
	bias=False,
	)

	# Frequency patch embedding (optional). Mirrors TransformerCNNEncoder2D:
	# rfft on a patch of length P gives P//2+1 complex values → 2*(P//2+1) reals,
	# which are projected to transformer_input_dim then concat-and-projected with
	# the time-domain patch embedding.
	self._use_freq_embed = bool(getattr(args, "use_freq_embed", False))
	# How the frequency branch encodes each patch spectrum.
	# "linear" (default, legacy): concat(real, imag) -> bias-free Linear.
	# NOTE: this makes the whole branch an exact linear reparameterisation
	# of `patch_embed`. rfft is a linear map, and patch_embed / freq_embed /
	# concat_proj are all bias-free with no nonlinearity between them, so
	# the fused output equals a single Linear(P -> D) of the raw patch.
	# It adds parameters, not expressivity. Kept as the default so that
	# existing `*_freq` checkpoints keep computing what they were trained to.
	# "spectral": concat(magnitude, phase) -> Linear -> GELU. \|F(x)\| is
	# nonlinear in x, so this genuinely enlarges the function class.
	# Shapes and parameter counts are identical, so a legacy checkpoint
	# will LOAD SILENTLY but be MISINTERPRETED -- only use this mode with
	# a model trained under it.
	self._freq_mode = str(getattr(args, "freq_mode", "linear")).lower()
	if self._freq_mode not in ("linear", "spectral"):
	raise ValueError(
	f"freq_mode must be 'linear' or 'spectral', got {self._freq_mode!r}"
	)
	if self._use_freq_embed:
	freq_input_dim = (patch_size // 2 + 1) * 2
	self.freq_embed = nn.Linear(freq_input_dim, transformer_input_dim, bias=False)
	self.concat_proj = nn.Linear(2 * transformer_input_dim, transformer_input_dim, bias=False)

	self.input_norm = nn.LayerNorm(transformer_input_dim)
	self.input_dropout = nn.Dropout(self.dropout)

	# Position-wise 2-layer MLP (same depth as requested).
	self.mlp = nn.Sequential(
	nn.Linear(transformer_input_dim, ff_dim),
	nn.GELU(),
	nn.Dropout(self.dropout),
	nn.Linear(ff_dim, transformer_input_dim),
	)
	self.out_norm = nn.LayerNorm(transformer_input_dim)

	if transformer_input_dim != output_dim:
	self.output_proj = nn.Linear(transformer_input_dim, output_dim)
	else:
	self.output_proj = nn.Identity()

	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"🔢 MLPEncoder2D: {total:,} params ({trainable:,} trainable)")

	# ------------------------------------------------------------------
	# Shared helpers
	# ------------------------------------------------------------------

	def _embed_and_pe(self, tok: torch.Tensor, chan_id: torch.Tensor, t_offset: int = 0) -> torch.Tensor:
	"""Add 2D sin-cos PE to a (C, T, D) patch-embed tensor."""
	C, T, D = tok.shape
	t_pos = torch.arange(t_offset, t_offset + T, device=tok.device, dtype=torch.float32)
	c_pos = chan_id.to(device=tok.device, dtype=torch.float32)

	use_only_time_pe = bool(getattr(self.args, "use_only_time_pe", False))
	use_only_channel_pe = bool(getattr(self.args, "use_only_channel_pe", False))

	if use_only_time_pe:
	pe = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
	pe = pe.unsqueeze(0).expand(C, -1, -1)
	elif use_only_channel_pe:
	pe = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
	pe = pe.unsqueeze(1).expand(-1, T, -1)
	else:
	if D % 4 == 0:
	pe_t = _get_1d_sincos_pos_embed(D // 2, t_pos, base=self.pe_base_t).to(tok.dtype)
	pe_c = _get_1d_sincos_pos_embed(D // 2, c_pos, base=self.pe_base_c).to(tok.dtype)
	pe = torch.cat(
	[pe_t.unsqueeze(0).expand(C, -1, -1),
	pe_c.unsqueeze(1).expand(-1, T, -1)],
	dim=-1,
	)
	else:
	pe_t = _get_1d_sincos_pos_embed(D, t_pos, base=self.pe_base_t).to(tok.dtype)
	pe_c = _get_1d_sincos_pos_embed(D, c_pos, base=self.pe_base_c).to(tok.dtype)
	pe = pe_t.unsqueeze(0).expand(C, -1, -1) + pe_c.unsqueeze(1).expand(-1, T, -1)

	return tok + pe

	# ------------------------------------------------------------------
	# Main forward
	# ------------------------------------------------------------------

	def forward(
	self,
	x: torch.Tensor,
	chan_id: torch.Tensor,
	attention_mask: torch.Tensor = None,
	) -> torch.Tensor:
	"""x: (C, L), chan_id: (C,), attention_mask: (C, T) or None. Returns (C, T, output_dim)."""
	C, L = x.shape
	if L % self.patch_size != 0:
	raise ValueError(f"L={L} not divisible by patch_size={self.patch_size}")

	T = L // self.patch_size
	tok = self.patch_embed(x.reshape(C, 1, L)).transpose(1, 2) # (C, T, D)

	if self._use_freq_embed:
	tok = self._fuse_freq(tok, x.reshape(C, T, self.patch_size))

	tok = self._embed_and_pe(tok, chan_id, t_offset=0)
	tok = self.input_norm(tok)
	tok = self.input_dropout(tok)

	# Position-wise MLP with residual.
	tok = tok + self.mlp(tok)
	tok = self.out_norm(tok)

	# Mask-out padding patches (zero their features) so downstream averaging
	# or projection doesn't leak padding content.
	if attention_mask is not None:
	m = attention_mask.to(dtype=tok.dtype, device=tok.device).unsqueeze(-1)
	tok = tok * m

	return self.output_proj(tok) # (C, T, output_dim)

	# ------------------------------------------------------------------
	# Compatibility stubs for callers that expect the transformer's
	# incremental API. MLP is position-wise so "incremental" is just a
	# single patch forward; we reuse `forward` for simplicity.
	# ------------------------------------------------------------------

	def tokenize_patches(
	self,
	x: torch.Tensor,
	chan_id: torch.Tensor,
	t_offset: int = 0,
	) -> torch.Tensor:
	if x.size(1) % self.patch_size != 0:
	raise ValueError(f"L={x.size(1)} not divisible by patch_size={self.patch_size}")
	C = x.size(0)
	T_new = x.size(1) // self.patch_size
	tok = self.patch_embed(x.reshape(C, 1, -1)).transpose(1, 2)

	if self._use_freq_embed:
	tok = self._fuse_freq(tok, x.reshape(C, T_new, self.patch_size))

	tok = self._embed_and_pe(tok, chan_id, t_offset=t_offset)
	return self.input_norm(tok)

	def _fuse_freq(self, tok: torch.Tensor, x_patches: torch.Tensor) -> torch.Tensor:
	"""Fuse a per-patch spectral embedding into the time-domain patch tokens.

	tok: (C, T, D) time-domain patch embeddings
	x_patches: (C, T, P) raw (z-normalised) patches
	Returns: (C, T, D)

	See the `_freq_mode` note in __init__ for why "linear" is expressively
	inert and "spectral" is not.
	"""
	x_fft = torch.fft.rfft(x_patches, dim=-1) # (C, T, P//2+1) complex
	if self._freq_mode == "spectral":
	# Magnitude is nonlinear in x; phase supplies the shift information
	# that magnitude alone discards.
	feats = torch.cat([x_fft.abs(), x_fft.angle()], dim=-1)
	freq_tok = torch.nn.functional.gelu(self.freq_embed(feats.to(tok.dtype)))
	else:
	feats = torch.cat([x_fft.real, x_fft.imag], dim=-1)
	freq_tok = self.freq_embed(feats.to(tok.dtype))
	return self.concat_proj(torch.cat([tok, freq_tok], dim=-1))