File size: 5,267 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import numpy as np
import torch
import torch.nn as nn
from models.utils import polyline_encoder
from einops import rearrange
import math
class SinusoidalPosEmb(nn.Module):
def __init__(self, dim, theta = 10000):
super().__init__()
self.dim = dim
self.theta = theta
def forward(self, x):
device = x.device
half_dim = self.dim // 2
emb = math.log(self.theta) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
emb = x[:, None] * emb[None, :]
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
return emb
class MTREncoder(nn.Module):
def __init__(self, config, use_pre_norm):
super().__init__()
self.model_cfg = config
dim = self.model_cfg.D_MODEL
# build polyline encoders
self.agent_polyline_encoder = self.build_polyline_encoder(
in_channels=self.model_cfg.NUM_INPUT_CONTEXT,
hidden_dim=self.model_cfg.NUM_CHANNEL_IN_MLP_AGENT,
num_layers=self.model_cfg.NUM_LAYER_IN_MLP_AGENT,
out_channels=dim
)
# Positional encoding
self.pos_encoding = nn.Sequential(
SinusoidalPosEmb(dim, theta = 10000),
nn.Linear(dim, dim),
nn.ReLU(),
nn.Linear(dim, dim)
)
self.team_one_query_embedding = nn.Embedding(1, dim)
self.team_two_query_embedding = nn.Embedding(1, dim)
self.ball_query_embedding = nn.Embedding(1, dim)
self.mlp_pe = nn.Sequential(
nn.Linear(2*dim, dim),
# nn.Linear(dim, dim),
nn.ReLU(),
nn.Linear(dim, dim)
)
# build transformer encoder layers
self.layer = nn.TransformerEncoderLayer(d_model=dim,
dropout=self.model_cfg.get('DROPOUT_OF_ATTN', 0.1),
nhead=self.model_cfg.NUM_ATTN_HEAD,
dim_feedforward=dim * 4,
norm_first=use_pre_norm,
batch_first=True)
self.transformer_encoder = nn.TransformerEncoder(self.layer, num_layers=self.model_cfg.NUM_ATTN_LAYERS)
self.num_out_channels = dim
### polyline encoder MLP
def build_polyline_encoder(self, in_channels, hidden_dim, num_layers, num_pre_layers=1, out_channels=None):
ret_polyline_encoder = polyline_encoder.PointNetPolylineEncoder(
in_channels=in_channels,
hidden_dim=hidden_dim,
num_layers=num_layers,
num_pre_layers=num_pre_layers,
out_channels=out_channels
)
return ret_polyline_encoder
def agent_query_embedding(self, index, num_agents=11, agent_type='sport'):
'''
Distinguish between team one, team two and ball (sport), or treat all
agents as a single "pedestrian" type (pedestrian). Works for any A.
- NBA: 5 + 5 + 1 = 11 (sport)
- Soccer / football: 11 + 11 + 1 = 23 (sport)
- SDD (target + neighbors): variable A (pedestrian)
'''
if agent_type == 'pedestrian':
# Single-type embedding (use team_one slot as "pedestrian") repeated A times.
ped_query = self.team_one_query_embedding(index)
return ped_query.repeat(num_agents, 1) # [A, D]
team_one_query = self.team_one_query_embedding(index)
team_two_query = self.team_two_query_embedding(index)
ball_query = self.ball_query_embedding(index)
team_size = (num_agents - 1) // 2
agent_query = torch.cat([
team_one_query.repeat(team_size, 1),
team_two_query.repeat(team_size, 1),
ball_query,
], dim=0)
return agent_query # [A, D]
def forward(self, past_traj, agent_type='sport', agent_mask=None):
"""
Args:
past_traj [B, A, T, 6]
agent_type 'sport' or 'pedestrian' (SDD).
agent_mask [B, A] bool; True = real, False = padded. If None, all real.
"""
past_traj_mask = torch.ones_like(past_traj[..., 0], dtype=torch.bool).to(past_traj.device)
obj_polylines_feature = self.agent_polyline_encoder(past_traj, past_traj_mask)
pos_encoding = self.pos_encoding(torch.arange(obj_polylines_feature.shape[1]).to(past_traj.device))
agent_query = self.agent_query_embedding(
torch.arange(1).to(past_traj.device),
num_agents=past_traj.shape[1],
agent_type=agent_type,
)
pos_encoding = self.mlp_pe(torch.cat([agent_query, pos_encoding], dim=-1))
obj_polylines_feature += pos_encoding.unsqueeze(0)
# Padding mask: src_key_padding_mask expects True where to IGNORE, so invert.
if agent_mask is not None:
kp_mask = ~agent_mask # [B, A] True=ignore
encoder_out = self.transformer_encoder(
obj_polylines_feature, src_key_padding_mask=kp_mask)
else:
encoder_out = self.transformer_encoder(obj_polylines_feature)
return encoder_out
|