refactor: move trainer/models/utils/config to tools/
This commit is contained in:
236
tools/models.py
Normal file
236
tools/models.py
Normal file
@@ -0,0 +1,236 @@
|
||||
"""Conditional Flow Matching model with Transformer backbone.
|
||||
|
||||
Predicts velocity field v_θ(x_t, t, cond) that transports noise to data.
|
||||
- x_t: (B, T, 3) noisy trajectory at time t
|
||||
- t: (B,) interpolation time ∈ [0,1]
|
||||
- cond: (B, 3) condition [dist_norm, log_dist, log_dur]
|
||||
- Output: (B, T, 3) velocity field
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sinusoidal time embedding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SinusoidalTimeEmbedding(nn.Module):
|
||||
"""Map scalar timestep t ∈ [0,1] to a d_model-dimensional vector."""
|
||||
|
||||
def __init__(self, d_model: int):
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
|
||||
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
t: (B,) scalar timesteps
|
||||
|
||||
Returns:
|
||||
(B, d_model) embedding vectors
|
||||
"""
|
||||
half = self.d_model // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(10000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half
|
||||
)
|
||||
args = t.unsqueeze(-1) * freqs.unsqueeze(0) # (B, half)
|
||||
return torch.cat([torch.sin(args), torch.cos(args)], dim=-1) # (B, d_model)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TrajectoryFlowModel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TrajectoryFlowModel(nn.Module):
|
||||
"""Conditional Flow Matching model for mouse trajectory generation.
|
||||
|
||||
Architecture:
|
||||
- SinusoidalTimeEmbedding: t scalar → d_model vector
|
||||
- Input projection: 3 → d_model
|
||||
- Learned positional embedding: (1, seq_len, d_model)
|
||||
- Time embed + Condition embed added as bias to all tokens
|
||||
- 4-layer TransformerEncoder (pre-norm, GELU, batch_first=True)
|
||||
- Output projection: d_model → 3
|
||||
|
||||
Args:
|
||||
seq_len: Number of time steps (default 64).
|
||||
d_model: Transformer hidden dimension (default 128).
|
||||
nhead: Number of attention heads (default 4).
|
||||
num_layers: Number of transformer layers (default 4).
|
||||
dim_feedforward: Feedforward hidden size (default 256).
|
||||
dropout: Dropout rate (default 0.1).
|
||||
cond_dim: Condition vector size (default 3).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
seq_len: int = 64,
|
||||
d_model: int = 128,
|
||||
nhead: int = 4,
|
||||
num_layers: int = 4,
|
||||
dim_feedforward: int = 256,
|
||||
dropout: float = 0.1,
|
||||
cond_dim: int = 3,
|
||||
):
|
||||
super().__init__()
|
||||
self.seq_len = seq_len
|
||||
self.d_model = d_model
|
||||
self.cond_dim = cond_dim
|
||||
|
||||
# Input projection: (forward, lateral, log_dt) → d_model
|
||||
self.input_proj = nn.Linear(3, d_model)
|
||||
|
||||
# Learned positional embedding
|
||||
self.pos_embed = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)
|
||||
|
||||
# Time embedding
|
||||
self.time_embed = SinusoidalTimeEmbedding(d_model)
|
||||
self.time_mlp = nn.Sequential(
|
||||
nn.Linear(d_model, d_model),
|
||||
nn.GELU(),
|
||||
nn.Linear(d_model, d_model),
|
||||
)
|
||||
|
||||
# Condition embedding
|
||||
self.cond_mlp = nn.Sequential(
|
||||
nn.Linear(cond_dim, d_model),
|
||||
nn.GELU(),
|
||||
nn.Linear(d_model, d_model),
|
||||
)
|
||||
|
||||
# Transformer encoder (pre-norm via norm_first=True)
|
||||
encoder_layer = nn.TransformerEncoderLayer(
|
||||
d_model=d_model,
|
||||
nhead=nhead,
|
||||
dim_feedforward=dim_feedforward,
|
||||
dropout=dropout,
|
||||
activation="gelu",
|
||||
batch_first=True,
|
||||
norm_first=True,
|
||||
)
|
||||
self.transformer = nn.TransformerEncoder(
|
||||
encoder_layer, num_layers=num_layers, enable_nested_tensor=False
|
||||
)
|
||||
|
||||
# Output projection: d_model → 3
|
||||
self.output_proj = nn.Linear(d_model, 3)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
t: torch.Tensor,
|
||||
cond: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Predict velocity field.
|
||||
|
||||
Args:
|
||||
x_t: (B, T, 3) noisy trajectory at interpolation time t
|
||||
t: (B,) interpolation timestep ∈ [0,1]
|
||||
cond: (B, 3) condition vector [dist_norm, log_dist, log_dur]
|
||||
|
||||
Returns:
|
||||
(B, T, 3) predicted velocity field
|
||||
"""
|
||||
# Project input tokens
|
||||
h = self.input_proj(x_t) # (B, T, d_model)
|
||||
|
||||
# Add positional embedding
|
||||
h = h + self.pos_embed # (B, T, d_model)
|
||||
|
||||
# Time embedding → bias added to all tokens
|
||||
t_emb = self.time_mlp(self.time_embed(t)) # (B, d_model)
|
||||
h = h + t_emb.unsqueeze(1) # broadcast over T
|
||||
|
||||
# Condition embedding → bias added to all tokens
|
||||
c_emb = self.cond_mlp(cond) # (B, d_model)
|
||||
h = h + c_emb.unsqueeze(1) # broadcast over T
|
||||
|
||||
# Transformer
|
||||
h = self.transformer(h) # (B, T, d_model)
|
||||
|
||||
# Output projection
|
||||
return self.output_proj(h) # (B, T, 3)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Legacy JointCVAE — kept for backward compatibility with generator.py
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class JointCVAE(nn.Module):
|
||||
"""Joint Conditional VAE for mouse trajectory generation (legacy).
|
||||
|
||||
Kept for backward compatibility with the existing generator.
|
||||
See TrajectoryFlowModel for the new approach.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
seq_len: int = 64,
|
||||
latent_dim: int = 32,
|
||||
hidden: int = 128,
|
||||
cond_dim: int = 3,
|
||||
):
|
||||
super().__init__()
|
||||
self.seq_len = seq_len
|
||||
self.latent_dim = latent_dim
|
||||
self.hidden = hidden
|
||||
self.cond_dim = cond_dim
|
||||
self.feat_dim = 3
|
||||
|
||||
self.enc_gru = nn.GRU(
|
||||
input_size=self.feat_dim + cond_dim,
|
||||
hidden_size=hidden,
|
||||
num_layers=2,
|
||||
batch_first=True,
|
||||
bidirectional=True,
|
||||
)
|
||||
self.enc_mu = nn.Linear(hidden * 2, latent_dim)
|
||||
self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
|
||||
|
||||
self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
|
||||
self.dec_gru = nn.GRU(
|
||||
input_size=latent_dim + cond_dim,
|
||||
hidden_size=hidden,
|
||||
num_layers=2,
|
||||
batch_first=True,
|
||||
)
|
||||
self.dec_out = nn.Linear(hidden, self.feat_dim)
|
||||
|
||||
def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
B, T, _ = seq.shape
|
||||
c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
|
||||
x_in = torch.cat([seq, c_exp], dim=-1)
|
||||
_, h_n = self.enc_gru(x_in)
|
||||
h_fwd = h_n[-2]
|
||||
h_bwd = h_n[-1]
|
||||
h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
|
||||
return self.enc_mu(h_cat), self.enc_logvar(h_cat)
|
||||
|
||||
def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
||||
B = z.shape[0]
|
||||
zc = torch.cat([z, cond], dim=-1)
|
||||
h0_flat = self.dec_h0(zc)
|
||||
h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
|
||||
inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
|
||||
out, _ = self.dec_gru(inp, h0)
|
||||
return self.dec_out(out)
|
||||
|
||||
def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
return Normal(mu, std).rsample()
|
||||
|
||||
def forward(
|
||||
self, seq: torch.Tensor, cond: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
mu, logvar = self.encode(seq, cond)
|
||||
z = self.reparameterise(mu, logvar)
|
||||
recon = self.decode(z, cond)
|
||||
return recon, mu, logvar
|
||||
Reference in New Issue
Block a user