File size: 1,798 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 | """
MotionTransformerGraphV5 — two-pass backbone with uncertainty-aware semantic
sparse graph (FutureInteractionGraphV5).
"""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v5 import FutureInteractionGraphV5
class MotionTransformerGraphV5(MotionTransformerGraph):
"""MotionTransformerGraph with uncertainty-aware semantic sparse 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):
super().__init__(model_config, logger, config,
graph_num_gnn_layers=graph_num_gnn_layers,
graph_dropout=graph_dropout)
D = self.dim
time_dim = D
self.future_graph = FutureInteractionGraphV5(
embed_dim = D,
future_steps = self.T_future,
num_agents = self.A,
num_heads = 4,
dropout = graph_dropout,
num_gnn_layers = graph_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,
)
params_graph = sum(p.numel() for p in self.future_graph.parameters())
logger.info("FutureInteractionGraphV5 parameters: {:,}".format(params_graph))
logger.info("Top-N neighbors: {:d} / {:d}".format(top_n_neighbors, self.A - 1))
|