| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class TemporalTrajectoryPredictor(nn.Module): |
| def __init__( |
| self, |
| num_tokens=512, |
| token_dim=768, |
| max_seq_len=64, |
| num_layers=12, |
| num_heads=8, |
| dim_feedforward=2048, |
| dropout=0.1 |
| ): |
| super().__init__() |
| self.num_tokens = num_tokens |
| self.token_dim = token_dim |
| self.max_seq_len = max_seq_len |
|
|
| self.time_pos_embed = nn.Parameter(torch.zeros(1, max_seq_len, token_dim)) |
| |
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=token_dim, |
| nhead=num_heads, |
| dim_feedforward=dim_feedforward, |
| dropout=dropout, |
| activation='gelu', |
| batch_first=True |
| ) |
| self.temporal_transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) |
| self.output_projection = nn.Linear(token_dim, token_dim) |
| |
| self._init_weights() |
|
|
| def _init_weights(self): |
| nn.init.normal_(self.time_pos_embed, std=0.02) |
| nn.init.xavier_uniform_(self.output_projection.weight) |
| if self.output_projection.bias is not None: |
| nn.init.constant_(self.output_projection.bias, 0) |
|
|
| def forward(self, history_frames, mask=None): |
| B, t, N, dim = history_frames.shape |
| |
| x = history_frames.transpose(1, 2).reshape(B * N, t, dim) |
| x = x + self.time_pos_embed[:, :t, :] |
| |
| if mask is None: |
| mask = torch.triu(torch.full((t, t), float('-inf'), device=x.device), diagonal=1) |
| |
| feat = self.temporal_transformer(x, mask=mask) |
| next_frames = self.output_projection(feat) |
| |
| return next_frames.reshape(B, N, t, dim).transpose(1, 2) |
|
|
|
|
| class SpatialGraphRefiner(nn.Module): |
| def __init__( |
| self, |
| num_tokens=512, |
| token_dim=768, |
| text_dim=1024, |
| max_graph_dist=10, |
| num_blocks=2, |
| num_heads=8, |
| dim_feedforward=2048, |
| dropout=0.1 |
| ): |
| super().__init__() |
| self.num_tokens = num_tokens |
| self.token_dim = token_dim |
| self.max_graph_dist = max_graph_dist |
| self.num_heads = num_heads |
| |
| distances = torch.arange(max_graph_dist + 1, dtype=torch.float32) |
| scales = torch.linspace(0.1, 1.0, num_heads) |
| self.register_buffer("graph_bias_table", -1.0 * (distances.unsqueeze(1) * scales.unsqueeze(0))) |
| |
| self.blocks = nn.ModuleList([ |
| nn.ModuleDict({ |
| "sa1": nn.MultiheadAttention(token_dim, num_heads, dropout=dropout, batch_first=True), |
| "ln1": nn.LayerNorm(token_dim), |
| "sa2": nn.MultiheadAttention(token_dim, num_heads, dropout=dropout, batch_first=True), |
| "ln2": nn.LayerNorm(token_dim), |
| "ca": nn.MultiheadAttention(token_dim, num_heads, kdim=text_dim, vdim=text_dim, dropout=dropout, batch_first=True), |
| "ln3": nn.LayerNorm(token_dim), |
| "ffn": nn.Sequential( |
| nn.Linear(token_dim, dim_feedforward), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(dim_feedforward, token_dim), |
| nn.Dropout(dropout) |
| ), |
| "ln4": nn.LayerNorm(token_dim) |
| }) for _ in range(num_blocks) |
| ]) |
|
|
| def _compute_graph_bias(self, shortest_path_matrix, B, t): |
| N = shortest_path_matrix.shape[-1] |
| dist_clipped = torch.clamp(shortest_path_matrix, 0, self.max_graph_dist).long() |
| bias = F.embedding(dist_clipped, self.graph_bias_table) |
| return bias.unsqueeze(1).expand(-1, t, -1, -1, -1).permute(0, 1, 4, 2, 3).reshape(B * t * self.num_heads, N, N) |
|
|
| def forward(self, next_frame_init, text_embed, shortest_path_matrix): |
| B, t, N, dim = next_frame_init.shape |
| graph_attn_bias = self._compute_graph_bias(shortest_path_matrix, B, t) |
| |
| x = next_frame_init.reshape(B * t, N, dim) |
| text_embed_expanded = text_embed.unsqueeze(1).expand(-1, t, -1, -1).reshape(B * t, -1, text_embed.shape[-1]) |
| |
| for b in self.blocks: |
| attn, _ = b["sa1"](x, x, x, attn_mask=graph_attn_bias) |
| x = b["ln1"](x + attn) |
| |
| attn, _ = b["sa2"](x, x, x, attn_mask=graph_attn_bias) |
| x = b["ln2"](x + attn) |
| |
| attn, _ = b["ca"](x, text_embed_expanded, text_embed_expanded) |
| x = b["ln3"](x + attn) |
| |
| x = b["ln4"](x + b["ffn"](x)) |
| |
| return x.reshape(B, t, N, dim) |
