File size: 2,097 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 | """
MotionTransformerGraphV6 — two-pass backbone with RAG-style sparse graph.
"""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v6 import FutureInteractionGraphV6
from models.interaction_baselines import build_interaction_module # E4: SRA_MODULE factory
class MotionTransformerGraphV6(MotionTransformerGraph):
"""MotionTransformerGraph with RAG-style sparse future interaction graph.
Extra constructor kwargs:
top_n_neighbors (int, default 5)
rel_traj_hidden (int, default 32)
y0_score_dim (int, default 32)
"""
def __init__(self, model_config, logger, config,
graph_num_gnn_layers: int = 2,
graph_dropout: float = 0.1,
top_n_neighbors: int = 5,
rel_traj_hidden: int = 32,
y0_score_dim: int = 32,
edge_mode: str = 'full',
neighbor_mode: str = 'rag'):
super().__init__(model_config, logger, config,
graph_num_gnn_layers=graph_num_gnn_layers,
graph_dropout=graph_dropout)
self.future_graph = build_interaction_module( # E4: sra|gameformer|c2f via env SRA_MODULE
embed_dim = self.dim,
future_steps = self.T_future,
num_agents = self.A,
num_heads = 4,
dropout = graph_dropout,
num_gnn_layers = graph_num_gnn_layers,
time_dim = self.dim,
top_n_neighbors = top_n_neighbors,
rel_traj_hidden = rel_traj_hidden,
y0_score_dim = y0_score_dim,
edge_mode = edge_mode,
neighbor_mode = neighbor_mode,
)
params_graph = sum(p.numel() for p in self.future_graph.parameters())
logger.info("FutureInteractionGraphV6 parameters: {:,}".format(params_graph))
logger.info("Top-N neighbors: {:d} / {:d} edge_mode: {} neighbor_mode: {}".format(
top_n_neighbors, self.A - 1, edge_mode, neighbor_mode))
|