diff --git a/src/velocity_prediction/train.py b/src/velocity_prediction/train.py index 0864526..40f6ac0 100644 --- a/src/velocity_prediction/train.py +++ b/src/velocity_prediction/train.py @@ -8,6 +8,7 @@ Usage: import argparse import os +import random import time import numpy as np from pathlib import Path @@ -39,9 +40,9 @@ def train_one_epoch( log_interval: int = 50, global_step: int = 0, use_amp: bool = True, - h_state: torch.Tensor = None, -) -> tuple[float, int, torch.Tensor]: - """Train for one epoch. Returns (avg_loss, updated_global_step, h_state_final).""" +) -> tuple[float, int]: + """Train for one epoch (stateless — each batch starts from zero hidden state). + Returns (avg_loss, updated_global_step).""" model.train() total_loss = 0.0 num_batches = 0 @@ -52,15 +53,11 @@ def train_one_epoch( tilt = batch["tilt"].to(device) # (B, S, 3) target = batch["v_body_target"].to(device) # (B, S, 2) - # Drop carried state if batch size changed (partial last batch / new scene) - if h_state is not None and h_state.shape[1] != events.shape[0]: - h_state = None - - # Per-step supervision over the whole sequence — TBPTT hidden state + # Stateless: each batch starts from zero hidden state with torch.amp.autocast(device.type, enabled=use_amp): - pred_seq, h_new = model(events, tilt, h_state) # (B, S, 2), (L, B, H) - loss_per_step = criterion(pred_seq, target) # (B, S, 2) - loss_per_step = loss_per_step.mean(-1) # (B, S) + pred_seq, _ = model(events, tilt, None) # (B, S, 2), h=None + loss_per_step = criterion(pred_seq, target) # (B, S, 2) + loss_per_step = loss_per_step.mean(-1) # (B, S) loss = loss_per_step.mean() optimizer.zero_grad() @@ -68,9 +65,6 @@ def train_one_epoch( scaler.step(optimizer) scaler.update() - # Detach hidden state for next batch — gradient only back to current seq_len - h_state = h_new.detach() - total_loss += loss.item() num_batches += 1 global_step += 1 @@ -87,7 +81,7 @@ def train_one_epoch( avg_loss = total_loss / max(num_batches, 1) print(f" Epoch {epoch} | Avg Loss: {avg_loss:.6f}") - return avg_loss, global_step, h_state + return avg_loss, global_step @torch.no_grad() @@ -150,15 +144,12 @@ def main(): print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)") print(f"AMP: {'enabled' if use_amp else 'disabled'}") - # Data loaders — TBPTT training requires strict temporal order - train_loader = create_tbptt_loader( - seq_len=train_cfg.seq_len, - batch_size=train_cfg.batch_size, - event_threshold=event_threshold, - event_use_log=train_cfg.event_use_log, - ) + # ── Randomised sequence lengths for curriculum ────────────── + SEQ_LEN_CHOICES = [64, 96, 128, 160, 192] + + # Validation uses a fixed seq_len for consistent comparison val_loader = create_val_loader( - seq_len=train_cfg.seq_len, + seq_len=128, stride=train_cfg.sliding_window_stride, batch_size=train_cfg.batch_size, num_workers=train_cfg.num_workers, @@ -216,20 +207,28 @@ def main(): ckpt_dir.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(log_dir=str(log_dir)) - print(f" seq_len={train_cfg.seq_len}, batch_size={train_cfg.batch_size}") + print(f" seq_len=random({min(SEQ_LEN_CHOICES)}~{max(SEQ_LEN_CHOICES)}), " + f"batch_size={train_cfg.batch_size}") print(f" lr={train_cfg.lr}, weight_decay={train_cfg.weight_decay}") print(f" log_dir={log_dir}, checkpoint_dir={ckpt_dir}\n") for epoch in range(start_epoch, train_cfg.epochs + 1): epoch_start = time.time() - # Reset GRU hidden state at epoch start — each epoch begins with h=0 - train_loss, global_step, _ = train_one_epoch( + # Random seq_len per epoch — GRU learns to handle varying temporal horizons + seq_len = random.choice(SEQ_LEN_CHOICES) + train_loader = create_tbptt_loader( + seq_len=seq_len, + batch_size=train_cfg.batch_size, + event_threshold=event_threshold, + event_use_log=train_cfg.event_use_log, + ) + + train_loss, global_step = train_one_epoch( model, train_loader, optimizer, criterion, scaler, device, epoch, writer, log_interval=train_cfg.log_interval, global_step=global_step, use_amp=use_amp, - h_state=None, ) val_loss = validate(model, val_loader, criterion, device, use_amp=use_amp) scheduler.step() @@ -238,6 +237,7 @@ def main(): current_lr = scheduler.get_last_lr()[0] print(f"Epoch {epoch:3d}/{train_cfg.epochs} | " + f"seq_len={seq_len:3d} | " f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | " f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")