feat: add stateful evaluation with model.step() for per-frame GRU inference

- Add VelocityPredictionModel.step() for single-frame forward pass
  with external GRU hidden state management
- Add evaluate_stateful() that processes frames sequentially,
  maintaining hidden state across timesteps (seq_len=1, stride=1, batch=1)
- Evaluation uses strict temporal ordering (num_workers=0)

Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
2026-06-21 20:58:57 +08:00
parent b0942180ba
commit 504332190d
2 changed files with 46 additions and 14 deletions
+27
View File
@@ -153,6 +153,33 @@ class VelocityPredictionModel(nn.Module):
v_body = self.head(last_hidden) # (B, 2)
return v_body
@torch.no_grad()
def step(self, events: torch.Tensor, tilt: torch.Tensor,
h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
"""
Single-frame stateful forward pass.
Processes one timestep and returns the prediction plus updated
GRU hidden state. Call repeatedly with the output ``h`` to
accumulate temporal context across frames.
Args:
events: (B, 1, 1, H, W) single-frame event map
tilt: (B, 1, 3) single-frame tilt vector
h: (num_layers, B, hidden_size) or None → zeros
Returns:
v_body: (B, 2) body-frame [v_right, v_forward]
h_new: (num_layers, B, hidden_size)
"""
cnn_feat = self.cnn(events) # (B, 1, 256)
pose_feat = self.pose_mlp(tilt) # (B, 1, 64)
fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, 1, 320)
_, h_new = self.gru(fused, h) # (num_layers, B, 128)
last_hidden = h_new[-1] # (B, 128)
v_body = self.head(last_hidden) # (B, 2)
return v_body, h_new
def count_parameters(model: nn.Module) -> int:
"""Count trainable parameters."""