first commit

engine/evaluate.py

@@ -0,0 +1,96 @@
+"""TuSimple-style accuracy/FP/FN evaluation (paper Eq. 10).
+
+Simplification flagged explicitly: the official TuSimple metric uses a 25px
+tolerance defined in original-resolution pixel space with per-clip point sampling.
+We don't retain clip structure (our data is single frames, not clips) or the
+per-sample letterbox scale at collation time, so we evaluate in normalized [0,1]
+canonical (640-wide) space with an equivalent threshold (25/640 ~= 0.039). This is a
+consistent proxy metric for comparing our own checkpoints, not a pixel-exact
+reproduction of the official script -- consistent with the val-split caveat already
+flagged in the project plan (no official test_label.json available).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from scipy.optimize import linear_sum_assignment
+
+from models.head import eval_curve
+
+X_THRESHOLD_NORM = 25.0 / 640.0
+CLS_PROB_THRESHOLD = 0.5
+
+
+@torch.no_grad()
+def evaluate(model, dataloader, device, sample_ys: torch.Tensor) -> dict:
+    model.eval()
+    total_correct_points = 0
+    total_gt_points = 0
+    total_fp_lanes = 0
+    total_fn_lanes = 0
+    total_pred_lanes = 0
+    total_gt_lanes = 0
+
+    for batch in dataloader:
+        images = batch["images"].to(device)
+        target_xs = batch["target_xs"].to(device)
+        target_valid_mask = batch["target_valid_mask"].to(device)
+        target_lane_valid = batch["target_lane_valid"].to(device)
+        sys_dev = sample_ys.to(device)
+
+        out = model(images)
+        probs = F.softmax(out["cls_logits"], dim=-1)[..., 1]  # (B, N)
+        pred_xs_all = eval_curve(out["curve_coeffs"], sys_dev.view(1, 1, -1).expand(*probs.shape, -1))  # (B, N, S)
+
+        B = images.shape[0]
+        for b in range(B):
+            pos_idx = (probs[b] > CLS_PROB_THRESHOLD).nonzero(as_tuple=True)[0]
+            gt_idx = target_lane_valid[b].nonzero(as_tuple=True)[0]
+
+            total_pred_lanes += pos_idx.numel()
+            total_gt_lanes += gt_idx.numel()
+
+            if gt_idx.numel() == 0:
+                total_fp_lanes += pos_idx.numel()
+                continue
+            if pos_idx.numel() == 0:
+                total_fn_lanes += gt_idx.numel()
+                gt_mask = target_valid_mask[b, gt_idx]
+                total_gt_points += int(gt_mask.sum().item())
+                continue
+
+            pred_xs = pred_xs_all[b, pos_idx]          # (P, S)
+            gt_xs = target_xs[b, gt_idx]                # (M, S)
+            gt_mask = target_valid_mask[b, gt_idx]      # (M, S)
+
+            diff = (pred_xs.unsqueeze(1) - gt_xs.unsqueeze(0)).abs()  # (P, M, S)
+            mask = gt_mask.unsqueeze(0).float()
+            denom = mask.sum(dim=-1).clamp(min=1.0)
+            cost = (diff * mask).sum(dim=-1) / denom    # (P, M)
+
+            pred_local, gt_local = linear_sum_assignment(cost.cpu().numpy())
+
+            matched_pred = set(pred_local.tolist())
+            matched_gt = set(gt_local.tolist())
+            total_fp_lanes += pos_idx.numel() - len(matched_pred)
+            total_fn_lanes += gt_idx.numel() - len(matched_gt)
+
+            for p_local, g_local in zip(pred_local, gt_local):
+                m = gt_mask[g_local]
+                if not m.any():
+                    continue
+                correct = (diff[p_local, g_local][m] < X_THRESHOLD_NORM).sum().item()
+                total_correct_points += correct
+                total_gt_points += int(m.sum().item())
+
+            unmatched_gt = set(range(gt_idx.numel())) - matched_gt
+            for g_local in unmatched_gt:
+                m = gt_mask[g_local]
+                total_gt_points += int(m.sum().item())
+
+    accuracy = total_correct_points / total_gt_points if total_gt_points > 0 else 0.0
+    fp = total_fp_lanes / total_pred_lanes if total_pred_lanes > 0 else 0.0
+    fn = total_fn_lanes / total_gt_lanes if total_gt_lanes > 0 else 0.0
+
+    model.train()
+    return {"accuracy": accuracy, "fp": fp, "fn": fn}

engine/train.py

@@ -0,0 +1,141 @@
+"""Phase 1 training loop: AdamW (discriminative LR), warmup+cosine, grad clipping,
+checkpointing, TuSimple accuracy/FP/FN eval per epoch. AMP/EMA are skipped for now
+since this runs on CPU (no usable CUDA on this machine -- see project notes); both
+are one-line additions once GPU training is available.
+"""
+from __future__ import annotations
+
+import math
+import os
+import time
+
+import torch
+from torch.utils.data import DataLoader
+
+from data.tusimple import TuSimpleMaskDataset, collate_fn
+from models.laneformer import LaneFormer
+from models.losses import compute_loss
+from engine.evaluate import evaluate
+
+
+def build_optimizer(model: LaneFormer, lr_backbone: float, lr_new: float, weight_decay: float) -> torch.optim.Optimizer:
+    backbone_params = list(model.backbone.parameters())
+    backbone_ids = {id(p) for p in backbone_params}
+    new_params = [p for p in model.parameters() if id(p) not in backbone_ids]
+    return torch.optim.AdamW(
+        [
+            {"params": backbone_params, "lr": lr_backbone},
+            {"params": new_params, "lr": lr_new},
+        ],
+        weight_decay=weight_decay,
+    )
+
+
+def build_scheduler(optimizer: torch.optim.Optimizer, warmup_steps: int, total_steps: int):
+    def lr_lambda(step: int) -> float:
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0)))
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train(cfg: dict) -> None:
+    torch.manual_seed(cfg["seed"])
+    torch.set_num_threads(cfg.get("num_threads", 12))
+    device = torch.device("cuda" if (cfg["train"].get("use_cuda") and torch.cuda.is_available()) else "cpu")
+    print(f"Using device: {device}", flush=True)
+
+    data_cfg = cfg["data"]
+    train_ds = TuSimpleMaskDataset(
+        root=data_cfg["tusimple_root"], split="train",
+        out_w=data_cfg["input_width"], out_h=data_cfg["input_height"],
+        max_lanes=data_cfg["max_lanes"], num_sample_ys=data_cfg["num_sample_ys"],
+        val_fraction=data_cfg["val_fraction"], seed=cfg["seed"],
+    )
+    val_ds = TuSimpleMaskDataset(
+        root=data_cfg["tusimple_root"], split="val",
+        out_w=data_cfg["input_width"], out_h=data_cfg["input_height"],
+        max_lanes=data_cfg["max_lanes"], num_sample_ys=data_cfg["num_sample_ys"],
+        val_fraction=data_cfg["val_fraction"], seed=cfg["seed"],
+    )
+    train_loader = DataLoader(
+        train_ds, batch_size=data_cfg["batch_size"], shuffle=True,
+        collate_fn=collate_fn, num_workers=data_cfg["num_workers"],
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=data_cfg["batch_size"], shuffle=False,
+        collate_fn=collate_fn, num_workers=data_cfg["num_workers"],
+    )
+    print(f"train={len(train_ds)} val={len(val_ds)} samples", flush=True)
+
+    model_cfg = cfg["model"]
+    model = LaneFormer(
+        backbone_name=model_cfg["backbone"], pretrained=model_cfg["pretrained"],
+        d_model=model_cfg["fusion_channels"], max_lanes=data_cfg["max_lanes"],
+        encoder_layers=model_cfg["encoder_layers"], decoder_layers=model_cfg["decoder_layers"],
+        nhead=model_cfg["attn_heads"], ffn_dim=model_cfg["ffn_dim"], dropout=model_cfg["dropout"],
+    ).to(device)
+
+    train_cfg = cfg["train"]
+    optimizer = build_optimizer(model, train_cfg["lr_backbone"], train_cfg["lr_new"], train_cfg["weight_decay"])
+    total_steps = train_cfg["epochs"] * len(train_loader)
+    scheduler = build_scheduler(optimizer, train_cfg["warmup_steps"], total_steps)
+
+    loss_cfg = cfg["loss"]
+    ckpt_dir = train_cfg["checkpoint_dir"]
+    os.makedirs(ckpt_dir, exist_ok=True)
+
+    best_acc = -1.0
+    global_step = 0
+
+    for epoch in range(train_cfg["epochs"]):
+        model.train()
+        epoch_start = time.time()
+        running = {"total": 0.0, "cls_loss": 0.0, "reg_loss": 0.0, "endpoint_loss": 0.0}
+        n_batches = 0
+
+        for batch in train_loader:
+            images = batch["images"].to(device)
+            targets = {k: batch[k].to(device) for k in
+                       ["target_xs", "target_valid_mask", "target_lane_valid", "target_endpoints"]}
+            sample_ys = batch["sample_ys"].to(device)
+
+            out = model(images)
+            loss_dict = compute_loss(
+                out, targets, sample_ys,
+                cls_weight=loss_cfg["cls_weight"], reg_weight=loss_cfg["reg_weight"],
+                endpoint_weight=loss_cfg["endpoint_weight"], bg_class_weight=loss_cfg["bg_class_weight"],
+            )
+
+            optimizer.zero_grad()
+            loss_dict["total"].backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg["grad_clip_norm"])
+            optimizer.step()
+            scheduler.step()
+
+            for k in running:
+                running[k] += float(loss_dict[k].detach())
+            n_batches += 1
+            global_step += 1
+
+            if global_step % train_cfg["log_every"] == 0:
+                avg = {k: v / n_batches for k, v in running.items()}
+                lr = scheduler.get_last_lr()[-1]
+                print(f"epoch {epoch} step {global_step} lr {lr:.2e} "
+                      f"loss {avg['total']:.4f} (cls {avg['cls_loss']:.4f} "
+                      f"reg {avg['reg_loss']:.4f} ep {avg['endpoint_loss']:.4f})", flush=True)
+
+        epoch_time = time.time() - epoch_start
+        metrics = evaluate(model, val_loader, device, val_ds[0]["sample_ys"] if len(val_ds) else train_ds[0]["sample_ys"])
+        print(f"[epoch {epoch}] time={epoch_time/60:.1f}min val_acc={metrics['accuracy']:.4f} "
+              f"fp={metrics['fp']:.4f} fn={metrics['fn']:.4f}", flush=True)
+
+        ckpt_path = os.path.join(ckpt_dir, "last.pt")
+        torch.save({"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, ckpt_path)
+        if metrics["accuracy"] > best_acc:
+            best_acc = metrics["accuracy"]
+            torch.save({"model": model.state_dict(), "epoch": epoch, "metrics": metrics},
+                       os.path.join(ckpt_dir, "best.pt"))
+            print(f"  new best (acc={best_acc:.4f}), saved best.pt", flush=True)