File size: 4,645 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 | """
FutureInteractionGraphV9b — V6 + receiver-adaptive readout (WITHOUT kinematic channels).
Tests the adaptive readout idea in isolation (3 channels like V6, not 8).
"""
import torch
import torch.nn as nn
from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
from models.graph_interaction_nba_v9 import AdaptiveRelTrajEncoder
from models.graph_interaction_nba_v5 import _heading_diff
class FutureInteractionGraphV9b(FutureInteractionGraphV6):
"""V6 scoring + adaptive temporal readout (3 channels)."""
def __init__(self, embed_dim, future_steps, num_agents,
num_heads=4, dropout=0.1, num_gnn_layers=2,
time_dim=128, top_n_neighbors=5, rel_traj_hidden=32,
y0_score_dim=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, top_n_neighbors=top_n_neighbors,
rel_traj_hidden=rel_traj_hidden, y0_score_dim=y0_score_dim)
# Replace with adaptive encoder (3 channels like V6)
self.rel_traj_encoder = AdaptiveRelTrajEncoder(
out_dim=embed_dim, T=future_steps,
D_hidden=rel_traj_hidden, num_heads=4,
in_channels=3, node_dim=embed_dim)
def forward(self, y_emb, y_abs, t_emb, tau, sigma_agent=None):
B, K, A, D = y_emb.shape
T = y_abs.shape[3]
N = self.top_n
# ---- Scoring (same as V6) ----
y0_flat = y_abs.reshape(B * K * A, T * 2)
y0_emb = self.y0_score_proj(y0_flat)
if sigma_agent is not None:
sigma_mean = sigma_agent.mean(dim=-1).reshape(B * K * A, 1)
tau_bka = sigma_mean.squeeze(-1)
else:
sigma_mean = torch.zeros(B * K * A, 1, device=y_abs.device)
tau_bka = (tau.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A).reshape(B * K * A))
node_feat = torch.cat([y0_emb, sigma_mean], dim=-1)
q_bka = self.W_q(node_feat)
k_bka = self.W_k(node_feat)
pos_bk = y_abs.reshape(B * K * A, T, 2)
edge_index_bk = self._make_batched_edge_index(B * K)
pos_i_t = pos_bk[edge_index_bk[1]]
pos_j_t = pos_bk[edge_index_bk[0]]
rel_pos_t = pos_j_t - pos_i_t
mean_rel = rel_pos_t.mean(dim=1)
std_rel = rel_pos_t.std(dim=1)
min_dist = rel_pos_t.norm(dim=-1).min(dim=1).values.unsqueeze(-1)
heading_full = _heading_diff(pos_i_t, pos_j_t)
heading_mean = heading_full.mean(dim=1)
q_i = q_bka[edge_index_bk[1]]
k_j = k_bka[edge_index_bk[0]]
semantic_score = (q_i * k_j).sum(dim=-1) * self.scale
geo_feat = torch.cat([mean_rel, std_rel, min_dist, heading_mean], dim=-1)
geo_bias = self.geo_mlp(geo_feat).squeeze(-1)
scores = semantic_score + geo_bias
# ---- Top-N ----
scores_grouped = scores.view(B * K * A, A - 1)
_, top_idx = scores_grouped.topk(N, dim=-1, sorted=False)
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)
# ---- Edge features: [rel_pos, heading_diff] + adaptive readout ----
heading_sparse = heading_full[mask_flat]
rel_pos_sparse = torch.cat([rel_pos_t[mask_flat], heading_sparse], dim=-1)
if sigma_agent is not None:
sigma_full = sigma_agent.reshape(B * K * A, T)
sigma_i_t = sigma_full[edge_index_bk[1][mask_flat]]
sigma_j_t = sigma_full[edge_index_bk[0][mask_flat]]
sigma_bias = sigma_i_t - sigma_j_t
else:
sigma_bias = None
# Receiver node embeddings for cross-attention
node_i_emb = y_emb.reshape(B * K * A, D)[edge_index_bk[1][mask_flat]]
edge_attr_sparse = self.rel_traj_encoder(
rel_pos_sparse, sigma_bias, node_i_emb=node_i_emb)
# ---- GNN + gated residual ----
edge_index_sparse = edge_index_bk[:, mask_flat]
temb_bka = (t_emb.unsqueeze(1).unsqueeze(2)
.expand(-1, K, A, -1).reshape(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)
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)
|