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:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user