| """MotionTransformerGraphV8 — V6 backbone with kinematic edge channels.""" |
| from models.backbone_graph import MotionTransformerGraph |
| from models.graph_interaction_nba_v8 import FutureInteractionGraphV8 |
|
|
|
|
| class MotionTransformerGraphV8(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 = FutureInteractionGraphV8( |
| 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"FutureInteractionGraphV8 (kinematic) parameters: {params:,}") |
|
|
