| """ |
| MotionTransformerGraphV13 — Graph as MODE SELECTOR. |
| |
| Instead of enriching y_emb (trajectory embedding), the graph adjusts |
| denoiser_cls (mode classification scores). The graph tells the model: |
| "given these agents' future interactions, mode K is more/less appropriate." |
| |
| Standard V6 graph enriches y_emb (same as before). |
| ADDITIONALLY, graph output produces a mode score adjustment that's added |
| to denoiser_cls. This directly influences which trajectory mode is selected. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| from einops import rearrange, repeat |
|
|
| from models.backbone_graph import MotionTransformerGraph |
| from models.graph_interaction_nba_v6 import FutureInteractionGraphV6 |
|
|
|
|
| class MotionTransformerGraphV13(MotionTransformerGraph): |
| """V6 graph + mode classification adjustment from graph output.""" |
|
|
| 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 = FutureInteractionGraphV6( |
| 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) |
|
|
| |
| self.mode_adjust = nn.Sequential( |
| nn.Linear(self.dim, self.dim // 2), |
| nn.ReLU(), |
| nn.Linear(self.dim // 2, 1), |
| ) |
| |
| nn.init.zeros_(self.mode_adjust[-1].weight) |
| nn.init.zeros_(self.mode_adjust[-1].bias) |
|
|
| p = sum(p.numel() for p in self.future_graph.parameters()) |
| logger.info(f"V13: Graph params: {p:,}, mode_adjust params: " |
| f"{sum(p.numel() for p in self.mode_adjust.parameters()):,}") |
|
|
| def _forward_impl(self, y, time, x_data, |
| y_0_for_graph=None, sigma_for_graph=None, |
| skip_graph=False): |
| if y.size(-1) == 2: |
| y = y.reshape((-1, self.model_cfg.NUM_PROPOSED_QUERY, |
| self.A, self.T_future * 2)) |
| device = y.device |
| B, K, A, _ = y.shape |
|
|
| encoder_out = self.context_encoder(x_data['past_traj_original_scale']) |
| encoder_out_batch = repeat(encoder_out, 'b a d -> b k a d', k=K, a=A) |
|
|
| y_emb = self.noisy_y_mlp(y) |
|
|
| time_ = time |
| if self.config.denoising_method == 'fm': |
| time = time * 1000.0 |
| t_emb = self.time_mlp(time) |
| t_emb_batch = repeat(t_emb, 'b d -> b k a d', b=B, k=K, a=A) |
|
|
| k_pe = self.motion_query_embedding( |
| torch.arange(self.model_cfg.NUM_PROPOSED_QUERY, device=device)) |
| k_pe_batch = repeat(k_pe, 'k d -> b k a d', b=B, a=A) |
| a_pe = self.agent_order_embedding( |
| torch.arange(self.model_cfg.CONTEXT_ENCODER.NUM_OF_ATTN_NEIGHBORS, device=device)) |
| a_pe_batch = repeat(a_pe, 'a d -> b k a d', b=B, k=K) |
|
|
| y_emb_k = rearrange(self.apply_PE(y_emb, k_pe_batch, a_pe_batch), 'b k a d -> (b a) k d') |
| y_emb_k = self.noisy_y_attn_k(y_emb_k) |
| y_emb = rearrange(y_emb_k, '(b a) k d -> b k a d', b=B, a=A) |
|
|
| y_emb_a = rearrange(y_emb, 'b k a d -> (b k) a d') |
| y_emb_a = self.noisy_y_attn_a(y_emb_a) |
| y_emb = rearrange(y_emb_a, '(b k) a d -> b k a d', b=B, k=K) |
|
|
| if self.training and self.config.get('drop_method', None) == 'emb': |
| m, k_drop = self.config.drop_logi_m, self.config.drop_logi_k |
| p_m = 1 / (1 + torch.exp(-k_drop * (time_ - m))) |
| p_m = p_m[:, None, None, None] |
| y_emb = y_emb.masked_fill(torch.rand_like(p_m) < p_m, 0.) |
|
|
| mode_cls_adjust = None |
| if not skip_graph: |
| y_graph_src = (y_0_for_graph.view(B, K, A, self.T_future, 2) |
| if y_0_for_graph is not None |
| else y.view(B, K, A, self.T_future, 2)) |
| y_graph_unnorm = self._unnormalize_y(y_graph_src) |
| init_pos = x_data['past_traj_original_scale'][:, :, -1, :2] |
| y_abs = y_graph_unnorm + init_pos.unsqueeze(1).unsqueeze(3) |
|
|
| tau = time_ |
| y_emb_graph = self.future_graph( |
| y_emb, y_abs, t_emb, tau, sigma_agent=sigma_for_graph) |
|
|
| |
| mode_cls_adjust = self.mode_adjust(y_emb_graph).squeeze(-1) |
|
|
| y_emb = y_emb_graph |
|
|
| emb_fusion = self.init_emb_fusion_mlp( |
| torch.cat((encoder_out_batch, y_emb, t_emb_batch), dim=-1)) |
| query_token = self.post_pe_cat_mlp( |
| self.apply_PE(emb_fusion, k_pe_batch, a_pe_batch)) |
| readout_token = self.motion_decoder(query_token, t_emb) |
|
|
| denoiser_x = self.reg_head(readout_token) |
| denoiser_cls = self.cls_head(readout_token).squeeze(-1) |
| logvar = self.logvar_head(readout_token) |
|
|
| |
| if mode_cls_adjust is not None: |
| denoiser_cls = denoiser_cls + mode_cls_adjust |
|
|
| return denoiser_x, denoiser_cls, logvar |
|
|
