File size: 1,796 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 | """
MotionTransformerGraphV2 — two-pass backbone with mode-specific graph edges.
Identical to MotionTransformerGraph (backbone_graph.py) except it uses
FutureInteractionGraphV2 instead of FutureInteractionGraph, so each of the
K denoising modes receives interaction edge features derived from its own
predicted future positions rather than a cross-mode average.
"""
from models.backbone_graph import MotionTransformerGraph
from models.graph_interaction_nba_v2 import FutureInteractionGraphV2
class MotionTransformerGraphV2(MotionTransformerGraph):
"""MotionTransformerGraph with mode-specific future interaction edges.
Constructor arguments are identical to MotionTransformerGraph.
Only __init__ is overridden to swap FutureInteractionGraph → V2.
All forward logic (_forward_impl, forward) is inherited unchanged.
"""
def __init__(self, model_config, logger, config,
graph_num_gnn_layers: int = 2,
graph_dropout: float = 0.1):
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 graph module with V2 (same hyperparams)
self.future_graph = FutureInteractionGraphV2(
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,
)
params_graph = sum(p.numel() for p in self.future_graph.parameters())
logger.info("FutureInteractionGraphV2 parameters: {:,}".format(params_graph))
|