| """MotionTransformerGraphV9b — V6 + adaptive readout (3ch, no kinematic).""" |
| from models.backbone_graph import MotionTransformerGraph |
| from models.graph_interaction_nba_v9b import FutureInteractionGraphV9b |
|
|
| class MotionTransformerGraphV9b(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 = FutureInteractionGraphV9b( |
| 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"FutureInteractionGraphV9b (adaptive readout) parameters: {params:,}") |
|
|
