File size: 1,806 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 | """
MotionTransformerGraphV3 — two-pass backbone with spatio-temporal graph edges.
Identical to MotionTransformerGraphV2 except it uses FutureInteractionGraphV3,
which encodes the full relative trajectory sequence (not just mean position)
using a small transformer (RelTrajEncoder) with self-attention over T timesteps.
"""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v3 import FutureInteractionGraphV3
class MotionTransformerGraphV3(MotionTransformerGraph):
"""MotionTransformerGraph with spatio-temporal future interaction edges.
Constructor arguments identical to MotionTransformerGraph, plus:
rel_traj_hidden (int, default 32): hidden dim inside RelTrajEncoder.
"""
def __init__(self, model_config, logger, config,
graph_num_gnn_layers: int = 2,
graph_dropout: float = 0.1,
rel_traj_hidden: 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
# Replace the V1/V2 graph module with V3 (spatio-temporal edge encoder)
self.future_graph = FutureInteractionGraphV3(
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,
rel_traj_hidden = rel_traj_hidden,
)
params_graph = sum(p.numel() for p in self.future_graph.parameters())
logger.info("FutureInteractionGraphV3 parameters: {:,}".format(params_graph))
|