张宗平
545802d62c
- src/tsmm/model/ts_encoder.py: Patchify (channel-independent unfold) + TSEncoder (shared per-channel patch linear, mean-pool over channels, learnable positional embedding, 2-layer pre-norm Transformer, LayerNorm). - configs/ts_encoder.yaml: encoder/projector hyperparameters. - tests/test_ts_encoder.py: 10 tests (shape, n_patches formula, window values, determinism, gradient flow, param sanity). - 61 tests passing; GPU fwd/bwd verified on RTX 3060 (39 MB peak). Spec clarifications (small tier, comet-build Step 4): - n_patches = (T-P)//S+1 ⇒ 127 for T=512 (design said 128; arithmetic slip). - 'Channel-independent' = shared per-channel patch embed + mean-pool → C-free output [B,n_patches,d]; channel identity carried by [属性] tokens. - Actual params ≈2.1M (design's '≈4M' was a rough estimate).
104 lines
3.8 KiB
Python
104 lines
3.8 KiB
Python
"""TS Encoder (T2.1): channel-independent patch embedding + lightweight Transformer.
|
|
|
|
Pipeline: ``[B, T, C]`` → Patchify → shared per-channel patch linear → mean-pool
|
|
over channels → +learnable positional embedding → 2-layer Transformer →
|
|
``[B, n_patches, d]``.
|
|
|
|
Design ref: ``docs/superpowers/specs/2026-06-29-ts-as-modality-design.md`` §2.
|
|
|
|
Spec clarifications (see ``tests/test_ts_encoder.py`` docstring):
|
|
- ``n_patches = (T - patch_len) // stride + 1`` (T=512,P=8,S=4 ⇒ 127).
|
|
- "Channel-independent" ⇒ shared per-channel patch embedding + mean-pool over
|
|
channels, so the output token sequence is C-free. Channel identity is supplied
|
|
separately by the ``[属性]`` text tokens in the multimodal splice (T2.3).
|
|
- Actual params ≈1.6M for the default compact config (design's "≈4M" was rough).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class Patchify(nn.Module):
|
|
"""Channel-independent patching along the time dimension.
|
|
|
|
Input : ``[B, T, C]``
|
|
Output : ``[B, n_patches, C, patch_len]`` where
|
|
``n_patches = (T - patch_len) // stride + 1``.
|
|
"""
|
|
|
|
def __init__(self, patch_len: int = 8, stride: int = 4) -> None:
|
|
super().__init__()
|
|
if patch_len <= 0 or stride <= 0:
|
|
raise ValueError("patch_len and stride must be positive")
|
|
self.patch_len = patch_len
|
|
self.stride = stride
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# x: [B, T, C] → unfold time into patches: [B, n_patches, C, patch_len]
|
|
return x.unfold(1, self.patch_len, self.stride)
|
|
|
|
def n_patches(self, T: int) -> int:
|
|
if T < self.patch_len:
|
|
raise ValueError(f"T={T} smaller than patch_len={self.patch_len}")
|
|
return (T - self.patch_len) // self.stride + 1
|
|
|
|
|
|
class TSEncoder(nn.Module):
|
|
"""TS Encoder: Patchify → shared patch linear (per channel) → channel pool
|
|
→ positional embedding → 2-layer Transformer.
|
|
|
|
Output is C-free: ``[B, n_patches, d]``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
d: int = 256,
|
|
layers: int = 2,
|
|
heads: int = 4,
|
|
patch_len: int = 8,
|
|
stride: int = 4,
|
|
ffn_mult: int = 4,
|
|
dropout: float = 0.1,
|
|
max_patches: int = 2048,
|
|
) -> None:
|
|
super().__init__()
|
|
if d % heads != 0:
|
|
raise ValueError(f"d ({d}) must be divisible by heads ({heads})")
|
|
self.d = d
|
|
self.patch_len = patch_len
|
|
self.stride = stride
|
|
self.patchify = Patchify(patch_len=patch_len, stride=stride)
|
|
|
|
# Shared per-channel patch embedding. Applied to each channel's patch
|
|
# independently, then we mean-pool over channels → C-free tokens.
|
|
self.patch_embed = nn.Linear(patch_len, d)
|
|
|
|
# Learnable positional embedding, large enough for any foreseeable T.
|
|
self.pos_embed = nn.Parameter(torch.zeros(max_patches, d))
|
|
nn.init.trunc_normal_(self.pos_embed, std=0.02)
|
|
|
|
encoder_layer = nn.TransformerEncoderLayer(
|
|
d_model=d,
|
|
nhead=heads,
|
|
dim_feedforward=d * ffn_mult,
|
|
dropout=dropout,
|
|
activation="gelu",
|
|
batch_first=True,
|
|
norm_first=True,
|
|
)
|
|
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=layers)
|
|
self.ln = nn.LayerNorm(d)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
# x: [B, T, C]
|
|
B, T, C = x.shape
|
|
patches = self.patchify(x) # [B, n_patches, C, patch_len]
|
|
n_patches = patches.shape[1]
|
|
|
|
h = self.patch_embed(patches) # [B, n_patches, C, d]
|
|
h = h.mean(dim=2) # [B, n_patches, d] (channel pool)
|
|
h = h + self.pos_embed[:n_patches] # add positional embedding
|
|
h = self.transformer(h) # [B, n_patches, d]
|
|
return self.ln(h)
|