File size: 1,165 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 | """MotionTransformerGraphV9 — kinematic + receiver-adaptive temporal readout."""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v9 import FutureInteractionGraphV9
class MotionTransformerGraphV9(MotionTransformerGraph):
def __init__(self, model_config, logger, config,
graph_num_gnn_layers=2, graph_dropout=0.1,
top_n_neighbors=5, rel_traj_hidden=32, y0_score_dim=32):
super().__init__(model_config, logger, config,
graph_num_gnn_layers=graph_num_gnn_layers,
graph_dropout=graph_dropout)
self.future_graph = FutureInteractionGraphV9(
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)
params = sum(p.numel() for p in self.future_graph.parameters())
logger.info(f"FutureInteractionGraphV9 (kinematic+adaptive) parameters: {params:,}")
|