File size: 6,823 Bytes
0d52562 | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """SDD dataloader for MoFlow with MULTI-SCENE BATCHING via per-batch padding.
Strategy:
- Sort scenes by A (number of agents) at init time.
- BatchSampler builds mini-batches from contiguous A-sorted indices so that
all scenes in a batch have similar A → padding waste is small (≤ 2×).
- Collate pads each batch to its own max_A and returns an `agent_mask`
[B, max_A] (True = real agent, False = padding) that the backbone and
graph module use to ignore padded slots.
Scenes within a batch attend ONLY to agents in their own scene via a
block-diagonal `traj_mask`; padded slots are zeroed out in the mask so
no information leaks across scenes or into padding.
"""
import os, pickle, math, random
import numpy as np
import torch
from torch.utils.data import Dataset, Sampler
from utils.normalization import normalize_min_max
def _sdd_to_scene(sdd, obs_len, pred_len):
"""Convert one raw scene [A, 20, 2] numpy array to torch tensors needed
by MoFlow: (past_6ch, fut_rel, fut_vel, init_pos).
Coords are already /50 and mean-centered per scene from preprocessing."""
traj = torch.from_numpy(sdd).float() # [A, 20, 2]
pre = traj[:, :obs_len].contiguous()
fut = traj[:, obs_len:].contiguous()
init = pre[:, -1:]
abs_ch = pre
rel_ch = pre - init
vel_ch = torch.cat([rel_ch[:, 1:] - rel_ch[:, :-1],
torch.zeros_like(rel_ch[:, -1:])], dim=1)
past_6 = torch.cat([abs_ch, rel_ch, vel_ch], dim=-1) # [A, T_past, 6]
fut_r = fut - init # [A, T_future, 2]
fut_vel = torch.cat([fut_r[:, 1:] - fut_r[:, :-1],
torch.zeros_like(fut_r[:, -1:])], dim=1)
return past_6, fut_r, fut_vel
class SortedByASampler(Sampler):
"""Yield single indices in A-sorted order with bucket-local shuffling.
Use with DataLoader(..., sampler=this, batch_size=B): consecutive indices
land in the same mini-batch so batches end up A-similar. Bucket size W =
4 * batch_size controls how much A-shuffling happens per epoch.
"""
def __init__(self, a_counts, batch_size, shuffle=True, seed=0):
self.a_counts = np.asarray(a_counts)
self.batch_size = int(batch_size)
self.shuffle = shuffle
self.seed = seed
self.epoch = 0
def set_epoch(self, epoch):
self.epoch = int(epoch)
def __iter__(self):
order = np.argsort(self.a_counts, kind='stable')
if self.shuffle:
rng = np.random.RandomState(self.seed + self.epoch)
W = 4 * self.batch_size
for start in range(0, len(order), W):
end = min(start + W, len(order))
rng.shuffle(order[start:end])
for idx in order.tolist():
yield int(idx)
def __len__(self):
return len(self.a_counts)
def collate_padded(batch):
"""Pad a list of `(past_6, fut_rel, fut_vel)` scenes (each [A_i, T, *]) to
a common max_A in this batch, stack, and emit an agent_mask.
Returns the same dict structure as dataloader_nba's seq_collate_nba plus
extra keys: 'agent_mask' [B, max_A] and 'A_per_scene' [B].
"""
B = len(batch)
A_list = [x[0].shape[0] for x in batch]
max_A = max(A_list)
T_past = batch[0][0].shape[1]
T_fut = batch[0][1].shape[1]
past_pad = torch.zeros(B, max_A, T_past, 6)
fut_pad = torch.zeros(B, max_A, T_fut, 2)
vel_pad = torch.zeros(B, max_A, T_fut, 2)
agent_mask = torch.zeros(B, max_A, dtype=torch.bool) # True = real
for i, (p6, fr, fv) in enumerate(batch):
A_i = p6.shape[0]
past_pad[i, :A_i] = p6
fut_pad[i, :A_i] = fr
vel_pad[i, :A_i] = fv
agent_mask[i, :A_i] = True
# Block-diagonal traj_mask zeroed on padding
N = B * max_A
traj_mask = torch.zeros(N, N)
flat_mask = agent_mask.view(-1) # [B*max_A]
for i in range(B):
s, e = i * max_A, (i + 1) * max_A
block = agent_mask[i].unsqueeze(0) & agent_mask[i].unsqueeze(1) # [max_A, max_A]
traj_mask[s:e, s:e] = block.float()
return {
'batch_size': torch.tensor(B),
'past_traj': past_pad, # [B, max_A, T_past, 6]
'fut_traj': fut_pad, # [B, max_A, T_fut, 2]
'past_traj_original_scale': past_pad.clone(), # same as past_traj (normalize afterwards if needed)
'fut_traj_original_scale': fut_pad.clone(),
'traj_mask': traj_mask, # [B*max_A, B*max_A]
'fut_traj_vel': vel_pad, # [B, max_A, T_fut, 2]
'agent_mask': agent_mask, # [B, max_A] True=real
}
class SDDDatasetMinMax(Dataset):
def __init__(
self,
obs_len=8, pred_len=12,
training=True,
cfg=None,
data_dir='',
split='train',
data_norm='min_max',
):
super().__init__()
self.obs_len = obs_len
self.pred_len = pred_len
self.cfg = cfg
self.data_norm = data_norm
self.training = training
path = os.path.join(data_dir, f'sdd_{split}_v2.pkl')
with open(path, 'rb') as f:
self.raw_scenes = pickle.load(f)
# Pre-compute A per scene for the BatchSampler
self.a_counts = [int(s.shape[0]) for s in self.raw_scenes]
print(f'[SDDDatasetMinMax] {split}: {len(self.raw_scenes)} scenes, '
f'A min/mean/max = {min(self.a_counts)}/{np.mean(self.a_counts):.1f}/{max(self.a_counts)}')
# Compute global min/max for min_max normalization (on training set)
if training and data_norm == 'min_max':
all_past, all_fut = [], []
for s in self.raw_scenes:
past_6, fut_r, _ = _sdd_to_scene(s, obs_len, pred_len)
all_past.append(past_6); all_fut.append(fut_r)
past_cat = torch.cat(all_past, dim=0)
fut_cat = torch.cat(all_fut, dim=0)
cfg.past_traj_max = past_cat.max()
cfg.past_traj_min = past_cat.min()
cfg.fut_traj_max = fut_cat.max()
cfg.fut_traj_min = fut_cat.min()
def __len__(self):
return len(self.raw_scenes)
def __getitem__(self, idx):
past_6, fut_r, fut_vel = _sdd_to_scene(self.raw_scenes[idx], self.obs_len, self.pred_len)
if self.data_norm == 'min_max':
past_6 = normalize_min_max(past_6, self.cfg.past_traj_min, self.cfg.past_traj_max, -1, 1).contiguous()
fut_r = normalize_min_max(fut_r, self.cfg.fut_traj_min, self.cfg.fut_traj_max, -1, 1).contiguous()
return past_6, fut_r, fut_vel
# Alias kept for backwards-compat with import lines in fm_sdd.py / fm_sdd_graph.py
seq_collate_sdd = collate_padded
|