File size: 10,190 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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | """
FutureInteractionGraphV3 β spatio-temporal edge encoding via RelTrajEncoder.
Difference from V2:
V2 encodes each directed edge as the mean relative future position
(averaged over T timesteps) β [E, 2] β MLP β [E, D].
V3 encodes the full relative trajectory sequence using a small
transformer (RelTrajEncoder): [E, T, 2] β [E, D].
Self-attention over T lets the model distinguish *when* and *how*
agent pairs interact, not just their average spatial offset.
Edge features are still computed per-mode (like V2), so each of the K
modes receives interaction context from its own predicted trajectories.
Memory:
Dominant intermediate: [B*K*E0, T, D_hidden] = [550K, 20, 32] β 1.4 GB.
Attention map: [B*K*E0, H, T, T] = [550K, 4, 20, 20] β 880 MB.
Total overhead ~3 GB beyond V2 β fits on an A6000.
"""
import torch
import torch.nn as nn
from models.graph_interaction_nba import FutureInteractionGraph
# ---------------------------------------------------------------------------
# Lightweight self-attention safe for large batch sizes (uses bmm not sdpa)
# ---------------------------------------------------------------------------
class _TemporalSelfAttn(nn.Module):
"""Pre-norm transformer block over a sequence dim, using torch.bmm.
nn.TransformerEncoderLayer uses scaled_dot_product_attention (flash attn)
which has a CUDA kernel batch-size limit (~2M). With E=550K edges and
H=4 heads the effective batch is 2.2M β exceeding the limit.
torch.bmm has no such restriction.
"""
def __init__(self, D: int, num_heads: int = 4):
super().__init__()
assert D % num_heads == 0
self.H = num_heads
self.Dh = D // num_heads
self.scale = self.Dh ** -0.5
self.qkv = nn.Linear(D, 3 * D, bias=False)
self.out = nn.Linear(D, D)
self.norm1 = nn.LayerNorm(D)
self.norm2 = nn.LayerNorm(D)
self.ff = nn.Sequential(
nn.Linear(D, D * 2), nn.ReLU(inplace=True), nn.Linear(D * 2, D),
)
def forward(self, x: torch.Tensor) -> torch.Tensor: # [E, T, D]
E, T, D = x.shape
H, Dh = self.H, self.Dh
# self-attention
res = x
x = self.norm1(x)
qkv = self.qkv(x).reshape(E, T, 3, H, Dh)
q = qkv[:, :, 0].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
k = qkv[:, :, 1].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
v = qkv[:, :, 2].permute(0, 2, 1, 3).reshape(E * H, T, Dh)
a = torch.bmm(q, k.transpose(1, 2)).mul_(self.scale).softmax(dim=-1)
out = torch.bmm(a, v).reshape(E, H, T, Dh).permute(0, 2, 1, 3).reshape(E, T, D)
x = self.out(out) + res
# feed-forward
x = x + self.ff(self.norm2(x))
return x
# ---------------------------------------------------------------------------
# Spatio-temporal relative trajectory encoder
# ---------------------------------------------------------------------------
class RelTrajEncoder(nn.Module):
"""Encode a relative trajectory sequence [E, T, 2] β [E, out_dim].
Architecture:
1. Per-timestep linear projection: [E, T, 2] β [E, T, D_hidden]
2. Add learned temporal position encoding
3. _TemporalSelfAttn (self-attention over T) β [E, T, D_hidden]
4. Attention-weighted pooling over T β [E, D_hidden]
5. Linear output projection β [E, out_dim]
The self-attention over T captures patterns like:
- when agents are closest (brief crossing vs sustained proximity)
- convergence vs divergence
- early vs late interaction in the prediction horizon
"""
def __init__(self, out_dim: int, T: int = 20,
D_hidden: int = 32, num_heads: int = 4,
in_channels: int = 2):
super().__init__()
self.input_proj = nn.Linear(in_channels, D_hidden)
self.t_pe = nn.Embedding(T, D_hidden) # learned temporal PE
self.attn = _TemporalSelfAttn(D_hidden, num_heads)
self.pool_w = nn.Linear(D_hidden, 1) # attention pooling weights
self.out_proj = nn.Linear(D_hidden, out_dim)
def forward(self, rel_pos_t: torch.Tensor,
sigma_bias: torch.Tensor = None) -> torch.Tensor:
"""
Args:
rel_pos_t: [E, T, 2] relative position at each future timestep
sigma_bias: [E, T] per-timestep uncertainty bias, or None.
Added to the pooling logits before softmax so that
timesteps where the target is uncertain and the source
is certain receive higher pooling weight.
Returns:
[E, out_dim]
"""
E, T, _ = rel_pos_t.shape
h = self.input_proj(rel_pos_t) # [E, T, D_hidden]
h = h + self.t_pe(torch.arange(T, device=h.device)) # temporal PE
h = self.attn(h) # [E, T, D_hidden]
w_logits = self.pool_w(h) # [E, T, 1]
if sigma_bias is not None:
w_logits = w_logits + sigma_bias.unsqueeze(-1) # [E, T, 1]
w = w_logits.softmax(dim=1) # [E, T, 1]
pooled = (h * w).sum(dim=1) # [E, D_hidden]
return self.out_proj(pooled) # [E, out_dim]
# ---------------------------------------------------------------------------
# V3 graph module
# ---------------------------------------------------------------------------
class FutureInteractionGraphV3(FutureInteractionGraph):
"""FutureInteractionGraph with spatio-temporal relative trajectory encoding.
Replaces the mean-position MLP (rel_pos_proj inherited from V1) with
RelTrajEncoder, which applies self-attention over the T future timesteps.
Edge features are computed per-mode (same as V2).
Extra constructor kwarg:
rel_traj_hidden (int, default 32): hidden dim inside 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,
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,
)
# Replace the V1 mean-position projection with the spatio-temporal encoder
del self.rel_pos_proj
self.rel_traj_encoder = RelTrajEncoder(
out_dim = embed_dim,
T = future_steps,
D_hidden = rel_traj_hidden,
num_heads = 4,
)
# ------------------------------------------------------------------
# Forward β per-mode edges + spatio-temporal encoding
# ------------------------------------------------------------------
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] 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 relative trajectory sequences [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]
# ---- Per-timestep uncertainty bias for RelTrajEncoder pooling ----
# sigma_agent [B, K, A, T]: per-agent per-timestep uncertainty.
# For edge (jβi): bias_t = Ο_i_t β Ο_j_t (uncertain target,
# certain source β higher pooling weight at that timestep).
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]] # [E, T] target
sigma_j_t = sigma_bka[edge_index_bk[0]] # [E, T] source
sigma_bias = sigma_i_t - sigma_j_t # [E, T]
# For GNN node_tau: mean over T (scalar per agent)
tau_bka = sigma_agent.mean(dim=-1).reshape(B * K * A)# [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_bk = self.rel_traj_encoder(rel_pos_t,
sigma_bias) # [B*K*E0, D]
# ---- Per-agent time embeddings ----------------------------------
temb_bka = (t_emb
.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A, -1)
.reshape(B * K * A, D)) # [B*K*A, D]
# ---- GNN pass ---------------------------------------------------
nodes = y_emb.reshape(B * K * A, D) # [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)) # [N, D]
out = orig + gate * self.out_proj(nodes) # [N, D]
return out.view(B, K, A, D)
|