feat: add TBPTT training with cross-batch hidden state carryover
- config: seq_len=128, batch_size=4 for long-sequence TBPTT - dataset: create_tbptt_loader with non-overlapping windows, strict temporal order - model: forward() accepts/exposes hidden state h; add step() for single-frame stateful inference - train: carry detached hidden state across batches, reset at epoch boundary - benchmark: fix model call for new (v_body, h) return signature Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
@@ -124,14 +124,17 @@ class VelocityPredictionModel(nn.Module):
|
||||
# nn.init.uniform_(self.head[-1].weight, -0.001, 0.001)
|
||||
# nn.init.zeros_(self.head[-1].bias)
|
||||
|
||||
def forward(self, events: torch.Tensor, tilt: torch.Tensor) -> torch.Tensor:
|
||||
def forward(self, events: torch.Tensor, tilt: torch.Tensor,
|
||||
h: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Args:
|
||||
events: (B, S, 1, H, W)
|
||||
tilt: (B, S, 3)
|
||||
h: (num_layers, B, hidden_size) or None → zeros (TBPTT state)
|
||||
|
||||
Returns:
|
||||
v_body: (B, S, 2) predicted body-frame [v_right, v_forward] at every timestep
|
||||
h_new: (num_layers, B, hidden_size) final GRU hidden state (detached by caller)
|
||||
"""
|
||||
B, S = events.shape[:2]
|
||||
|
||||
@@ -145,12 +148,12 @@ class VelocityPredictionModel(nn.Module):
|
||||
# Fuse per frame
|
||||
fused = torch.cat([cnn_feat, pose_feat], dim=-1) # (B, S, 320)
|
||||
|
||||
# GRU temporal modelling
|
||||
gru_out, _ = self.gru(fused) # (B, S, 128)
|
||||
# GRU temporal modelling — accepts external hidden state for TBPTT
|
||||
gru_out, h_new = self.gru(fused, h) # (B, S, 128), (num_layers, B, 128)
|
||||
|
||||
# Head regression over every timestep (single-layer GRU → gru_out == h_n[-1] at t=-1)
|
||||
v_body = self.head(gru_out) # (B, S, 128) → (B, S, 2)
|
||||
return v_body
|
||||
return v_body, h_new
|
||||
|
||||
@torch.no_grad()
|
||||
def step(self, events: torch.Tensor, tilt: torch.Tensor,
|
||||
@@ -195,7 +198,7 @@ if __name__ == "__main__":
|
||||
B, S, H, W = 4, 8, 240, 320
|
||||
events = torch.randn(B, S, 1, H, W)
|
||||
tilt = torch.randn(B, S, 3)
|
||||
out = model(events, tilt)
|
||||
out, h = model(events, tilt)
|
||||
print(f"Input events: {events.shape}")
|
||||
print(f"Input tilt: {tilt.shape}")
|
||||
print(f"Output: {out.shape} (should be [4, 8, 2])")
|
||||
|
||||
Reference in New Issue
Block a user