Huang Qi
4d414fd180
- Add .gitignore for Python/data/models - Add matplotlib>=3.8.0 for eval plots - Add PretrainConfig, FinetuneConfig, BalabitAdapterConfig, EvalConfig dataclasses
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""Tests for TrajectoryFlowModel architecture."""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import pytest
|
|
|
|
from ai_mouse.models import TrajectoryFlowModel
|
|
|
|
|
|
class TestTrajectoryFlowModel:
|
|
"""Test the Conditional Flow Matching model."""
|
|
|
|
@pytest.fixture
|
|
def model(self):
|
|
return TrajectoryFlowModel(
|
|
seq_len=64, d_model=128, nhead=4, num_layers=4,
|
|
dim_feedforward=256, dropout=0.1, cond_dim=3,
|
|
)
|
|
|
|
def test_output_shape(self, model):
|
|
"""(4, 64, 3) input → (4, 64, 3) output."""
|
|
batch = 4
|
|
x_t = torch.randn(batch, 64, 3)
|
|
t = torch.rand(batch)
|
|
cond = torch.randn(batch, 3)
|
|
out = model(x_t, t, cond)
|
|
assert out.shape == (batch, 64, 3)
|
|
|
|
def test_single_sample(self, model):
|
|
"""(1, 64, 3) works."""
|
|
x_t = torch.randn(1, 64, 3)
|
|
t = torch.rand(1)
|
|
cond = torch.randn(1, 3)
|
|
out = model(x_t, t, cond)
|
|
assert out.shape == (1, 64, 3)
|
|
|
|
def test_deterministic(self, model):
|
|
"""Eval mode, same input → same output."""
|
|
model.eval()
|
|
x_t = torch.randn(2, 64, 3)
|
|
t = torch.tensor([0.3, 0.7])
|
|
cond = torch.randn(2, 3)
|
|
with torch.no_grad():
|
|
out1 = model(x_t, t, cond)
|
|
out2 = model(x_t, t, cond)
|
|
torch.testing.assert_close(out1, out2)
|
|
|
|
def test_different_timesteps(self, model):
|
|
"""t=0.1 vs t=0.9 gives different output."""
|
|
model.eval()
|
|
x_t = torch.randn(1, 64, 3)
|
|
cond = torch.randn(1, 3)
|
|
with torch.no_grad():
|
|
out_early = model(x_t, torch.tensor([0.1]), cond)
|
|
out_late = model(x_t, torch.tensor([0.9]), cond)
|
|
assert not torch.allclose(out_early, out_late, atol=1e-5)
|
|
|
|
def test_gradient_flows(self, model):
|
|
"""Backward works, grad on x_t exists."""
|
|
model.train()
|
|
x_t = torch.randn(2, 64, 3, requires_grad=True)
|
|
t = torch.rand(2)
|
|
cond = torch.randn(2, 3)
|
|
out = model(x_t, t, cond)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
assert x_t.grad is not None
|
|
assert x_t.grad.shape == (2, 64, 3)
|
|
assert x_t.grad.abs().sum() > 0
|