张宗平
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).
87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
"""Tests for TS Encoder (T2.1).
|
||
|
||
Spec clarification notes (recorded in tasks.md / commit):
|
||
- n_patches = (T - patch_len) // stride + 1 → T=512,P=8,S=4 ⇒ 127 (design's
|
||
"128" was an arithmetic slip; downstream is channel-wise so 1-token diff is nil).
|
||
- "Channel-independent" = shared per-channel patch embedding + mean-pool over
|
||
channels → output [B, n_patches, d]. Channel identity carried by [属性] text tokens.
|
||
- Actual params ≈1.6M for a compact 2-layer d=256 encoder (design's "≈4M" was rough).
|
||
"""
|
||
import torch
|
||
|
||
from tsmm.model.ts_encoder import Patchify, TSEncoder
|
||
|
||
|
||
class TestPatchify:
|
||
def test_basic_shape(self):
|
||
pf = Patchify(patch_len=8, stride=4)
|
||
x = torch.randn(2, 512, 5)
|
||
patches = pf(x)
|
||
assert patches.shape == (2, 127, 5, 8)
|
||
|
||
def test_n_patches_formula(self):
|
||
pf = Patchify(patch_len=8, stride=4)
|
||
for T, expected in [(512, 127), (128, 31), (256, 63), (16, 3)]:
|
||
x = torch.randn(1, T, 3)
|
||
assert pf(x).shape[1] == expected, f"T={T}"
|
||
|
||
def test_non_overlapping(self):
|
||
# stride == patch_len ⇒ non-overlapping
|
||
pf = Patchify(patch_len=8, stride=8)
|
||
x = torch.randn(1, 512, 4)
|
||
assert pf(x).shape == (1, 64, 4, 8)
|
||
|
||
def test_values_are_correct_windows(self):
|
||
# patch i covers x[:, i*stride : i*stride+patch_len, :]
|
||
pf = Patchify(patch_len=4, stride=2)
|
||
x = torch.arange(12, dtype=torch.float).view(1, 12, 1)
|
||
patches = pf(x) # [1, n_patches=5, 1, 4]
|
||
# patch 0 = x[0:4]
|
||
assert torch.allclose(patches[0, 0, 0], torch.tensor([0.0, 1.0, 2.0, 3.0]))
|
||
# patch 1 = x[2:6]
|
||
assert torch.allclose(patches[0, 1, 0], torch.tensor([2.0, 3.0, 4.0, 5.0]))
|
||
|
||
|
||
class TestTSEncoder:
|
||
def test_output_shape(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4)
|
||
x = torch.randn(2, 512, 5)
|
||
out = enc(x)
|
||
assert out.shape == (2, 127, 256)
|
||
|
||
def test_different_T(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4)
|
||
for T, n in [(128, 31), (256, 63), (512, 127)]:
|
||
out = enc(torch.randn(1, T, 3))
|
||
assert out.shape == (1, n, 256)
|
||
|
||
def test_different_C(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4)
|
||
for C in [1, 5, 17]:
|
||
out = enc(torch.randn(1, 512, C))
|
||
assert out.shape == (1, 127, 256) # output C-free
|
||
|
||
def test_param_count_reasonable(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4)
|
||
n_params = sum(p.numel() for p in enc.parameters())
|
||
# compact 2-layer d=256 encoder ≈ 1.6M; sanity range 0.5M–6M
|
||
assert 500_000 < n_params < 6_000_000, f"params={n_params}"
|
||
|
||
def test_deterministic_in_eval(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4).eval()
|
||
x = torch.randn(1, 128, 3)
|
||
with torch.no_grad():
|
||
o1 = enc(x)
|
||
o2 = enc(x)
|
||
assert torch.allclose(o1, o2)
|
||
|
||
def test_gradient_flows(self):
|
||
enc = TSEncoder(d=256, layers=2, heads=4, patch_len=8, stride=4)
|
||
x = torch.randn(1, 128, 3, requires_grad=False)
|
||
out = enc(x)
|
||
loss = out.pow(2).sum()
|
||
loss.backward()
|
||
grads = [p.grad for p in enc.parameters() if p.grad is not None]
|
||
assert all(g is not None for g in grads)
|
||
assert any(g.abs().sum() > 0 for g in grads)
|