File size: 3,748 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 | """
FutureInteractionGraphV2 — mode-specific edge features.
Difference from V1 (graph_interaction_nba.py):
V1 averages the K predicted trajectories into one representative trajectory
per scene, computes edge features from that average, and then replicates them
across all K modes. Every mode therefore sees identical inter-agent geometry.
V2 computes edge features *separately for each mode*, so agent i's message
from agent j in mode k reflects that mode's specific predicted positions,
not the cross-mode mean.
Architecture and GNN code (FutureGeomAttnLayer) are unchanged — only the
edge-feature-construction step in forward() is different.
Memory note:
V1 computed edges at B scale (27.5 K edges) then replicated.
V2 computes edges directly at B*K scale (550 K edges) in one step.
The GNN forward pass size is identical in both versions (55 K nodes,
550 K edges); only the edge-feature tensor construction is larger in V2
(B*K*E0*T*2 ≈ 88 MB extra intermediate before projection — negligible
on an A6000).
"""
import torch
from models.graph_interaction_nba import FutureInteractionGraph
class FutureInteractionGraphV2(FutureInteractionGraph):
"""FutureInteractionGraph with per-mode (mode-specific) edge features.
Constructor arguments are identical to FutureInteractionGraph.
Only forward() is overridden.
"""
def forward(
self,
y_emb: torch.Tensor, # [B, K, A, D]
y_abs: torch.Tensor, # [B, K, A, T, 2] absolute-ish future pos
t_emb: torch.Tensor, # [B, D] time embedding from backbone
tau: torch.Tensor, # [B] ∈ [0, 1]
sigma_agent: torch.Tensor = None,# [B, K, A] per-agent uncertainty, 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
# ---- Per-mode edge features (B*K scenes) [B*K*E0, D] -------------
# Flatten all B*K mode-trajectories together, then compute relative
# positions per directed edge for each mode independently.
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]
# Mean relative position over future timesteps for each (mode, edge)
rel_pos_bk = (pos_bk[edge_index_bk[0]] -
pos_bk[edge_index_bk[1]]).mean(dim=1) # [B*K*E0, 2]
edge_attr_bk = self.rel_pos_proj(rel_pos_bk) # [B*K*E0, D]
# ---- Per-agent time / uncertainty embeddings [B*K*A, D] ----------
temb_bka = (t_emb
.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A, -1)
.reshape(B * K * A, D))
if sigma_agent is not None:
s = sigma_agent.mean(dim=-1) if sigma_agent.dim() == 4 else sigma_agent
tau_bka = s.reshape(B * K * A)
else:
tau_bka = (tau
.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A)
.reshape(B * K * A))
# ---- GNN pass on B*K*A nodes -------------------------------------
nodes = y_emb.reshape(B * K * A, D)
for layer in self.gnn_layers:
nodes = layer(nodes, edge_index_bk, edge_attr_bk,
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)
|