File size: 7,327 Bytes
d4cbafd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """
FutureInteractionGraphV4 — learnable sparse graph with spatio-temporal encoding.
Two-stage design:
Stage 1 (cheap, all edges):
Compute multi-scale geometric features [mean_rel, std_rel, min_dist] → [E, 5]
Score each edge with a small MLP → [E] scalar importance.
Select top-N neighbors per agent (hard, per-mode).
Stage 2 (expensive, selected edges only):
Run RelTrajEncoder (self-attention over T) only on the N*B*K*A selected edges.
GNN runs on the resulting sparse graph.
Philosophy:
- Per-mode: scoring and encoding both use per-mode predicted positions.
- Timestep-specific: RelTrajEncoder retains full T-step temporal structure.
- Speed: RelTrajEncoder runs on B*K*A*N edges instead of B*K*A*(A-1).
With N=5, A=11 → ~2x fewer encoder calls; with N=3 → ~3x fewer.
Gradient flow:
The scorer MLP receives gradients through the selected edges' GNN outputs.
No straight-through needed — the scorer naturally learns to rank important
neighbors because good selections lead to better trajectory predictions.
"""
import torch
import torch.nn as nn
from models.graph_interaction_nba import FutureInteractionGraph
from models.graph_interaction_nba_v3 import RelTrajEncoder
# ---------------------------------------------------------------------------
# V4 graph module
# ---------------------------------------------------------------------------
class FutureInteractionGraphV4(FutureInteractionGraph):
"""Sparse future interaction graph with learnable top-N neighbor selection.
Extra constructor kwargs (beyond FutureInteractionGraph):
top_n_neighbors (int, default 5): neighbors to keep per agent per mode.
rel_traj_hidden (int, default 32): hidden dim in RelTrajEncoder.
"""
def __init__(self, embed_dim: int, future_steps: int, num_agents: int,
num_heads: int = 4, dropout: float = 0.1,
num_gnn_layers: int = 2, time_dim: int = 128,
top_n_neighbors: int = 5, rel_traj_hidden: int = 32):
super().__init__(
embed_dim = embed_dim,
future_steps = future_steps,
num_agents = num_agents,
num_heads = num_heads,
dropout = dropout,
num_gnn_layers = num_gnn_layers,
time_dim = time_dim,
)
assert top_n_neighbors < num_agents, \
f"top_n_neighbors ({top_n_neighbors}) must be < num_agents ({num_agents})"
self.top_n = top_n_neighbors
# Stage 1: cheap edge scorer
# Features: mean_rel[2] + std_rel[2] + min_dist[1] = 5-dim
# Learnable: MLP learns which geometric patterns indicate important interactions.
self.edge_scorer = nn.Sequential(
nn.Linear(5, 16),
nn.ReLU(inplace=True),
nn.Linear(16, 1),
)
# Stage 2: full spatio-temporal encoder (selected edges only)
del self.rel_pos_proj # replace V1's mean-position MLP
self.rel_traj_encoder = RelTrajEncoder(
out_dim = embed_dim,
T = future_steps,
D_hidden = rel_traj_hidden,
num_heads = 4,
)
# ------------------------------------------------------------------
# Forward
# ------------------------------------------------------------------
def forward(
self,
y_emb: torch.Tensor, # [B, K, A, D]
y_abs: torch.Tensor, # [B, K, A, T, 2]
t_emb: torch.Tensor, # [B, D]
tau: torch.Tensor, # [B] ∈ [0, 1]
sigma_agent: torch.Tensor = None, # [B, K, A, T] or None
) -> torch.Tensor: # [B, K, A, D]
B, K, A, D = y_emb.shape
T = y_abs.shape[3]
E0 = self._E0 # A*(A-1), edges per scene
N = self.top_n
# ---- Full per-mode relative trajectories [B*K*E0, T, 2] ----------
pos_bk = y_abs.reshape(B * K * A, T, 2) # [B*K*A, T, 2]
edge_index_bk = self._make_batched_edge_index(B * K) # [2, B*K*E0]
rel_pos_t = (pos_bk[edge_index_bk[0]] -
pos_bk[edge_index_bk[1]]) # [B*K*E0, T, 2]
# ---- Stage 1: cheap scorer on all edges --------------------------
mean_rel = rel_pos_t.mean(dim=1) # [E, 2]
std_rel = rel_pos_t.std(dim=1) # [E, 2]
min_dist = (rel_pos_t.norm(dim=-1)
.min(dim=1).values
.unsqueeze(-1)) # [E, 1]
score_feat = torch.cat([mean_rel, std_rel, min_dist], dim=-1) # [E, 5]
scores = self.edge_scorer(score_feat).squeeze(-1) # [B*K*E0]
# ---- Top-N selection per target agent ----------------------------
# Edge ordering from _make_single_edge_index: outer loop is target i,
# so each target has exactly (A-1) contiguous edges → safe to view.
scores_grouped = scores.view(B * K * A, A - 1) # [B*K*A, A-1]
_, top_idx = scores_grouped.topk(N, dim=-1, sorted=False)# [B*K*A, N]
mask = torch.zeros(B * K * A, A - 1,
device=scores.device, dtype=torch.bool)
mask.scatter_(1, top_idx, True)
mask_flat = mask.view(-1) # [B*K*E0] bool
# ---- Stage 2: RelTrajEncoder on selected edges only --------------
rel_pos_sparse = rel_pos_t[mask_flat] # [B*K*A*N, T, 2]
if sigma_agent is not None:
sigma_bka = sigma_agent.reshape(B * K * A, T) # [B*K*A, T]
sigma_i_t = sigma_bka[edge_index_bk[1][mask_flat]]# [B*K*A*N, T]
sigma_j_t = sigma_bka[edge_index_bk[0][mask_flat]]# [B*K*A*N, T]
sigma_bias = sigma_i_t - sigma_j_t # [B*K*A*N, T]
tau_bka = sigma_agent.mean(dim=-1).reshape(B * K * A)
else:
sigma_bias = None
tau_bka = (tau
.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A)
.reshape(B * K * A)) # [B*K*A]
edge_attr_sparse = self.rel_traj_encoder(
rel_pos_sparse, sigma_bias
) # [B*K*A*N, D]
# ---- Sparse GNN pass ---------------------------------------------
edge_index_sparse = edge_index_bk[:, mask_flat] # [2, B*K*A*N]
temb_bka = (t_emb
.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A, -1)
.reshape(B * K * A, D)) # [B*K*A, D]
nodes = y_emb.reshape(B * K * A, D)
for layer in self.gnn_layers:
nodes = layer(nodes, edge_index_sparse, edge_attr_sparse,
temb_agent=temb_bka, tau=tau_bka)
# ---- Gated residual ----------------------------------------------
orig = y_emb.reshape(B * K * A, D)
gate = self.gate_proj(torch.cat([orig, nodes], dim=-1))
out = orig + gate * self.out_proj(nodes)
return out.view(B, K, A, D)
|