diff --git a/src/velocity_prediction/evaluate.py b/src/velocity_prediction/evaluate.py index a0d790f..6a74bc2 100644 --- a/src/velocity_prediction/evaluate.py +++ b/src/velocity_prediction/evaluate.py @@ -19,33 +19,37 @@ from src.velocity_prediction.config import train_cfg, VELOCITY_MEAN, VELOCITY_ST @torch.no_grad() -def evaluate( +def evaluate_stateful( model: nn.Module, loader, device: torch.device, ) -> dict: """ - Run evaluation on a dataloader. + Stateful evaluation: process frames one-by-one, maintaining GRU + hidden state across timesteps for full temporal context. + + The loader must yield single-frame samples (seq_len=1, stride=1) + in strict temporal order (num_workers=0). Returns: dict with keys: - preds: np.ndarray (N, 2) predicted [vx, vy] - targets: np.ndarray (N, 2) ground truth [vx, vy] + preds: np.ndarray (N, 2) predicted [v_right, v_forward] + targets: np.ndarray (N, 2) ground truth """ model.eval() all_preds = [] all_targets = [] + h = None for batch in loader: - events = batch["events"].to(device) - tilt = batch["tilt"].to(device) - target = batch["v_body_target"].to(device) # (B, S, 2) + events = batch["events"].to(device) # (B=1, S=1, 1, H, W) + tilt = batch["tilt"].to(device) # (B=1, S=1, 3) + target = batch["v_body_target"].to(device) # (B=1, S=1, 2) - pred = model(events, tilt) # (B, 2) - target_last = target[:, -1, :] # (B, 2) + pred, h = model.step(events, tilt, h) # (1, 2), (1, 1, 128) all_preds.append(pred.cpu().numpy()) - all_targets.append(target_last.cpu().numpy()) + all_targets.append(target[:, -1, :].cpu().numpy()) preds = np.concatenate(all_preds, axis=0) targets = np.concatenate(all_targets, axis=0) @@ -178,13 +182,14 @@ def main(): for scene in TEST_SCENES: loader = create_val_loader( scene_names=[scene], - seq_len=train_cfg.seq_len, - batch_size=train_cfg.batch_size, - num_workers=2, + seq_len=1, + stride=1, + batch_size=1, + num_workers=0, # strict temporal order event_threshold=train_cfg.event_threshold, event_use_log=train_cfg.event_use_log, ) - results = evaluate(model, loader, device) + results = evaluate_stateful(model, loader, device) n = len(results["preds"]) print(f" [{scene}] RMSE vx={results['rmse_x']:.4f} vy={results['rmse_y']:.4f} " f"xy={results['rmse_xy']:.4f} samples={n}") diff --git a/src/velocity_prediction/model.py b/src/velocity_prediction/model.py index 196cc00..b21f514 100644 --- a/src/velocity_prediction/model.py +++ b/src/velocity_prediction/model.py @@ -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."""