mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 21:13:20 +08:00
99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
"""Set Transformer 模块自检脚本
|
||
|
||
约定:
|
||
- 每个 test_xxx 函数返回 None；失败时直接抛 AssertionError 或 RuntimeError。
|
||
- main() 收集所有 test_xxx 函数依次执行，统计 pass/fail 并以非零状态退出。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
import traceback
|
||
from typing import Callable, List, Tuple
|
||
|
||
import torch
|
||
|
||
from lnp_ml.modeling.layers.set_transformer import SetTransformer
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
# Tests
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
|
||
def test_shape_preserved() -> None:
|
||
st = SetTransformer(d_model=256, num_heads=8, n_layers=4)
|
||
x = torch.randn(2, 4, 256)
|
||
out = st(x)
|
||
assert out.shape == (2, 4, 256), f"out shape {out.shape} != (2, 4, 256)"
|
||
|
||
|
||
def test_variable_set_size() -> None:
|
||
"""n_chem=3（无 MPNN）和 4（有 MPNN）都要支持。"""
|
||
st = SetTransformer(d_model=64, num_heads=8, n_layers=2)
|
||
for n in (3, 4):
|
||
x = torch.randn(5, n, 64)
|
||
assert st(x).shape == (5, n, 64), f"n={n} 形状不符"
|
||
|
||
|
||
def test_permutation_equivariance() -> None:
|
||
"""无位置编码，输出应随输入 token 置换而同步置换。"""
|
||
torch.manual_seed(0)
|
||
st = SetTransformer(d_model=64, num_heads=8, n_layers=3).eval()
|
||
x = torch.randn(2, 4, 64)
|
||
perm = torch.tensor([2, 0, 3, 1])
|
||
out = st(x)
|
||
out_perm = st(x[:, perm, :])
|
||
assert torch.allclose(out[:, perm, :], out_perm, atol=1e-4), "未满足置换等变性"
|
||
|
||
|
||
def test_grad_flows() -> None:
|
||
st = SetTransformer(d_model=32, num_heads=4, n_layers=2)
|
||
x = torch.randn(3, 4, 32, requires_grad=True)
|
||
st(x).sum().backward()
|
||
assert x.grad is not None and torch.isfinite(x.grad).all(), "梯度异常"
|
||
|
||
|
||
def test_isab_shape() -> None:
|
||
st = SetTransformer(d_model=64, num_heads=8, n_layers=2, block="isab", num_inducing=8)
|
||
x = torch.randn(2, 4, 64)
|
||
assert st(x).shape == (2, 4, 64), "ISAB 形状不符"
|
||
|
||
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
# Runner
|
||
# ──────────────────────────────────────────────────────────────────────
|
||
|
||
def _collect_tests() -> List[Tuple[str, Callable[[], None]]]:
|
||
return [
|
||
(name, fn) for name, fn in globals().items()
|
||
if name.startswith("test_") and callable(fn)
|
||
]
|
||
|
||
|
||
def main() -> int:
|
||
tests = _collect_tests()
|
||
pass_count = 0
|
||
fail: List[Tuple[str, str]] = []
|
||
|
||
for name, fn in tests:
|
||
try:
|
||
fn()
|
||
except Exception:
|
||
fail.append((name, traceback.format_exc()))
|
||
print(f"[FAIL] {name}")
|
||
else:
|
||
pass_count += 1
|
||
print(f"[ OK ] {name}")
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"Passed: {pass_count}/{len(tests)}")
|
||
if fail:
|
||
print(f"\n{'='*60}\nFailures:")
|
||
for name, tb in fail:
|
||
print(f"\n--- {name} ---\n{tb}")
|
||
return 1
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main()) |