53 lines
2.2 KiB
Python
53 lines
2.2 KiB
Python
"""Per-query prediction heads: lane/background classification + cubic curve regression.
|
|
|
|
Phase 1 note: curve coefficients are regressed directly (unconstrained) rather than
|
|
with LSTR's full numerical-reparameterization trick -- kept simple for the first
|
|
trainable pass; revisit if training shows y^3-term gradient instability.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class MLP(nn.Module):
|
|
def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int = 2):
|
|
super().__init__()
|
|
dims = [in_dim] + [hidden_dim] * (num_layers - 1) + [out_dim]
|
|
layers = []
|
|
for i in range(len(dims) - 1):
|
|
layers.append(nn.Linear(dims[i], dims[i + 1]))
|
|
if i < len(dims) - 2:
|
|
layers.append(nn.ReLU(inplace=True))
|
|
self.net = nn.Sequential(*layers)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.net(x)
|
|
|
|
|
|
class PredictionHeads(nn.Module):
|
|
"""x = k*y^3 + m*y^2 + n*y + b, plus (y_start, y_end), plus lane/background logits."""
|
|
|
|
def __init__(self, d_model: int, hidden_dim: int = 128):
|
|
super().__init__()
|
|
self.cls_head = MLP(d_model, hidden_dim, 2, num_layers=2)
|
|
self.curve_head = MLP(d_model, hidden_dim, 4, num_layers=3) # k, m, n, b
|
|
self.endpoint_head = MLP(d_model, hidden_dim, 2, num_layers=2) # y_start, y_end (pre-sigmoid)
|
|
|
|
def forward(self, queries: torch.Tensor) -> dict:
|
|
"""queries: (B, N_lanes, d_model)"""
|
|
cls_logits = self.cls_head(queries) # (B, N, 2)
|
|
curve_coeffs = self.curve_head(queries) # (B, N, 4)
|
|
endpoints = torch.sigmoid(self.endpoint_head(queries)) # (B, N, 2) in [0,1]
|
|
return {"cls_logits": cls_logits, "curve_coeffs": curve_coeffs, "endpoints": endpoints}
|
|
|
|
|
|
def eval_curve(curve_coeffs: torch.Tensor, ys: torch.Tensor) -> torch.Tensor:
|
|
"""curve_coeffs: (..., 4) = [k,m,n,b]; ys: (...,S) or (S,) broadcastable normalized y.
|
|
Returns x(y) with shape (..., S).
|
|
"""
|
|
k, m, n, b = curve_coeffs.unbind(dim=-1) # each (...,)
|
|
k, m, n, b = k.unsqueeze(-1), m.unsqueeze(-1), n.unsqueeze(-1), b.unsqueeze(-1)
|
|
y = ys
|
|
return k * y**3 + m * y**2 + n * y + b
|