hexone2086
56aa10a503
- model outputs (B, S, 2) instead of (B, 2) — GRU output at every timestep - train/val loss computed over all S timesteps with reduction=none - benchmark/evaluate.py takes pred[:, -1, :] for final-step evaluation - added per-step loss logging (8 evenly spaced steps) to TensorBoard Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
278 lines
10 KiB
Python
278 lines
10 KiB
Python
"""
|
|
Training loop for VelocityPredictionModel.
|
|
|
|
Usage:
|
|
python -m src.velocity_prediction.train [--device cuda:0]
|
|
python -m src.velocity_prediction.train --resume checkpoints/epoch_020_val_0.123456.pt
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
|
|
from src.velocity_prediction.config import train_cfg, model_cfg
|
|
from src.velocity_prediction.model import VelocityPredictionModel, count_parameters
|
|
from src.velocity_prediction.dataset import create_train_loader, create_val_loader
|
|
|
|
|
|
def set_seed(seed: int):
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def train_one_epoch(
|
|
model: nn.Module,
|
|
loader,
|
|
optimizer: torch.optim.Optimizer,
|
|
criterion: nn.Module,
|
|
scaler: torch.cuda.amp.GradScaler,
|
|
device: torch.device,
|
|
epoch: int,
|
|
writer: SummaryWriter,
|
|
log_interval: int = 50,
|
|
global_step: int = 0,
|
|
use_amp: bool = True,
|
|
) -> tuple[float, int]:
|
|
"""Train for one epoch. Returns (avg_loss, updated_global_step)."""
|
|
model.train()
|
|
total_loss = 0.0
|
|
num_batches = 0
|
|
start_time = time.time()
|
|
|
|
for batch_idx, batch in enumerate(loader):
|
|
events = batch["events"].to(device) # (B, S, 1, H, W)
|
|
tilt = batch["tilt"].to(device) # (B, S, 3)
|
|
target = batch["v_body_target"].to(device) # (B, S, 2)
|
|
|
|
# Per-step supervision over the whole sequence
|
|
with torch.amp.autocast(device.type, enabled=use_amp):
|
|
pred_seq = model(events, tilt) # (B, S, 2)
|
|
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()
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
|
|
total_loss += loss.item()
|
|
num_batches += 1
|
|
global_step += 1
|
|
|
|
if batch_idx % log_interval == 0:
|
|
elapsed = time.time() - start_time
|
|
print(f" Epoch {epoch} | Batch {batch_idx} | Loss: {loss.item():.6f} | {elapsed:.1f}s")
|
|
writer.add_scalar("train/loss_batch", loss.item(), global_step)
|
|
# Per-step loss quantiles — observe GRU stabilization across timesteps
|
|
S = loss_per_step.shape[1]
|
|
for i in range(0, S, max(1, S // 8)):
|
|
writer.add_scalar(f"train/loss_step_{i:02d}",
|
|
loss_per_step[:, i].mean().item(), global_step)
|
|
|
|
avg_loss = total_loss / max(num_batches, 1)
|
|
print(f" Epoch {epoch} | Avg Loss: {avg_loss:.6f}")
|
|
return avg_loss, global_step
|
|
|
|
|
|
@torch.no_grad()
|
|
def validate(
|
|
model: nn.Module,
|
|
loader,
|
|
criterion: nn.Module,
|
|
device: torch.device,
|
|
use_amp: bool = True,
|
|
) -> float:
|
|
"""Validate. Returns average loss."""
|
|
model.eval()
|
|
total_loss = 0.0
|
|
num_batches = 0
|
|
|
|
for batch in loader:
|
|
events = batch["events"].to(device)
|
|
tilt = batch["tilt"].to(device)
|
|
target = batch["v_body_target"].to(device)
|
|
|
|
with torch.amp.autocast(device.type, enabled=use_amp):
|
|
pred_seq = model(events, tilt) # (B, S, 2)
|
|
loss_per_step = criterion(pred_seq, target) # (B, S, 2)
|
|
loss = loss_per_step.mean(-1).mean()
|
|
|
|
total_loss += loss.item()
|
|
num_batches += 1
|
|
|
|
avg_loss = total_loss / max(num_batches, 1)
|
|
return avg_loss
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--device", type=str, default="cuda",
|
|
help="CUDA device, e.g. 'cuda:0', 'cuda:1' (default: 'cuda')")
|
|
parser.add_argument("--resume", type=str, default=None,
|
|
help="Path to checkpoint .pt file to resume training from")
|
|
parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True,
|
|
help="Enable Automatic Mixed Precision (default: True)")
|
|
parser.add_argument("--threshold", type=float, default=None,
|
|
help="Event brightness change threshold (overrides config, default: same as train_cfg)")
|
|
args = parser.parse_args()
|
|
use_amp = args.amp
|
|
|
|
set_seed(train_cfg.seed)
|
|
device = torch.device(args.device if torch.cuda.is_available() and "cuda" in args.device else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
# Override threshold from CLI if provided
|
|
event_threshold = args.threshold if args.threshold is not None else train_cfg.event_threshold
|
|
if args.threshold is not None:
|
|
print(f"Event threshold overridden: {train_cfg.event_threshold} → {event_threshold}")
|
|
|
|
# Create model
|
|
model = VelocityPredictionModel()
|
|
model.to(device)
|
|
scaler = torch.amp.GradScaler(device.type, enabled=use_amp)
|
|
total_params = count_parameters(model)
|
|
print(f"Model parameters: {total_params:,} ({total_params/1e6:.3f} M)")
|
|
print(f"AMP: {'enabled' if use_amp else 'disabled'}")
|
|
|
|
# Data loaders
|
|
train_loader = create_train_loader(
|
|
seq_len=train_cfg.seq_len,
|
|
stride=train_cfg.sliding_window_stride,
|
|
batch_size=train_cfg.batch_size,
|
|
num_workers=train_cfg.num_workers,
|
|
event_threshold=event_threshold,
|
|
event_use_log=train_cfg.event_use_log,
|
|
)
|
|
val_loader = create_val_loader(
|
|
seq_len=train_cfg.seq_len,
|
|
stride=train_cfg.sliding_window_stride,
|
|
batch_size=train_cfg.batch_size,
|
|
num_workers=train_cfg.num_workers,
|
|
event_threshold=event_threshold,
|
|
event_use_log=train_cfg.event_use_log,
|
|
)
|
|
|
|
# Optimizer & scheduler
|
|
optimizer = torch.optim.AdamW(
|
|
model.parameters(),
|
|
lr=train_cfg.lr,
|
|
weight_decay=train_cfg.weight_decay,
|
|
)
|
|
scheduler = torch.optim.lr_scheduler.StepLR(
|
|
optimizer,
|
|
step_size=train_cfg.lr_scheduler_step,
|
|
gamma=train_cfg.lr_scheduler_gamma,
|
|
)
|
|
# criterion = nn.SmoothL1Loss(reduction='none')
|
|
criterion = nn.MSELoss(reduction='none')
|
|
|
|
# ── Resume from checkpoint ────────────────────────────────────
|
|
start_epoch = 1
|
|
global_step = 0
|
|
best_val_loss = float("inf")
|
|
run_id = None
|
|
|
|
if args.resume is not None:
|
|
ckpt_path = Path(args.resume)
|
|
if not ckpt_path.exists():
|
|
raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
|
|
ckpt = torch.load(ckpt_path, map_location="cpu")
|
|
|
|
model.load_state_dict(ckpt["model_state_dict"])
|
|
optimizer.load_state_dict(ckpt["optimizer_state_dict"])
|
|
if "scheduler_state_dict" in ckpt:
|
|
scheduler.load_state_dict(ckpt["scheduler_state_dict"])
|
|
start_epoch = ckpt.get("epoch", 0) + 1
|
|
global_step = ckpt.get("global_step", 0)
|
|
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
|
run_id = ckpt.get("run_id", None)
|
|
|
|
print(f"Resumed from checkpoint: {ckpt_path}")
|
|
print(f" Resumed epoch={ckpt.get('epoch', '?')}, global_step={global_step}, "
|
|
f"best_val_loss={best_val_loss:.6f}")
|
|
else:
|
|
print(f"\nStarting training for {train_cfg.epochs} epochs...")
|
|
|
|
# Logging & checkpoint — run-specific subdirectories for isolation + resume continuity
|
|
if run_id is None:
|
|
run_id = time.strftime("run_%Y%m%d_%H%M%S")
|
|
log_dir = Path(train_cfg.log_dir) / run_id
|
|
log_dir.mkdir(parents=True, exist_ok=True)
|
|
ckpt_dir = Path(train_cfg.checkpoint_dir) / run_id
|
|
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" 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()
|
|
|
|
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,
|
|
)
|
|
val_loss = validate(model, val_loader, criterion, device, use_amp=use_amp)
|
|
scheduler.step()
|
|
|
|
epoch_time = time.time() - epoch_start
|
|
current_lr = scheduler.get_last_lr()[0]
|
|
|
|
print(f"Epoch {epoch:3d}/{train_cfg.epochs} | "
|
|
f"Train Loss: {train_loss:.6f} | Val Loss: {val_loss:.6f} | "
|
|
f"LR: {current_lr:.2e} | Time: {epoch_time:.1f}s")
|
|
|
|
writer.add_scalar("train/loss_epoch", train_loss, epoch)
|
|
writer.add_scalar("val/loss", val_loss, epoch)
|
|
writer.add_scalar("lr", current_lr, epoch)
|
|
|
|
# Save checkpoint
|
|
if epoch % train_cfg.save_interval == 0:
|
|
ckpt_path = ckpt_dir / f"epoch_{epoch:03d}_val_{val_loss:.6f}.pt"
|
|
torch.save({
|
|
"epoch": epoch,
|
|
"model_state_dict": model.state_dict(),
|
|
"optimizer_state_dict": optimizer.state_dict(),
|
|
"scheduler_state_dict": scheduler.state_dict(),
|
|
"global_step": global_step,
|
|
"best_val_loss": best_val_loss,
|
|
"run_id": run_id,
|
|
"val_loss": val_loss,
|
|
}, ckpt_path)
|
|
print(f" Checkpoint saved: {ckpt_path}")
|
|
|
|
# Save best model
|
|
if val_loss < best_val_loss:
|
|
best_val_loss = val_loss
|
|
best_path = ckpt_dir / "best.pt"
|
|
torch.save({
|
|
"epoch": epoch,
|
|
"model_state_dict": model.state_dict(),
|
|
"optimizer_state_dict": optimizer.state_dict(),
|
|
"scheduler_state_dict": scheduler.state_dict(),
|
|
"global_step": global_step,
|
|
"best_val_loss": best_val_loss,
|
|
"run_id": run_id,
|
|
"val_loss": val_loss,
|
|
}, best_path)
|
|
print(f" Best model updated: {best_path} (val_loss={val_loss:.6f})")
|
|
|
|
writer.close()
|
|
print(f"\nTraining complete. Best val loss: {best_val_loss:.6f}")
|
|
print(f"Best checkpoint: {ckpt_dir / 'best.pt'}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|