File size: 7,000 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 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 | """
Re-evaluate an LED SDD checkpoint using the MID evaluation protocol:
per-pedestrian full-horizon (12-frame) ADE / final-frame FDE,
best_of_20 per pedestrian, averaged across all target pedestrians,
×50 to report in pixels.
Usage:
python eval_sdd_led_mid_protocol.py --exp baseline_v2 --epoch 60
python eval_sdd_led_mid_protocol.py --exp graph_sigma_v2 --epoch 60 --use_graph
"""
import argparse, os, sys, random, torch, numpy as np
from torch.utils.data import DataLoader
from data.dataloader_sdd import SDDDataset, sdd_seq_collate
from models.model_led_initializer import LEDInitializer as InitializationModel
from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel
from trainer.train_sdd_led import NUM_Tau
from utils.config import Config
def build_models(cfg, use_graph, use_v6_graph, ckpt_path, device):
model = CoreDenoisingModel(past_len=cfg.past_frames).to(device)
core_ckpt = torch.load(cfg.pretrained_core_denoising_model, map_location='cpu')
model.load_state_dict(core_ckpt['model_dict'])
model.eval()
init = InitializationModel(
t_h=cfg.past_frames, d_h=6,
t_f=cfg.future_frames, d_f=2, k_pred=20).to(device)
ckpt = torch.load(ckpt_path, map_location='cpu')
init.load_state_dict(ckpt['model_initializer_dict'])
init.eval()
graph = None
if use_graph:
from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper
graph = FutureInteractionGraphV6Wrapper(
num_agents=64, future_steps=cfg.future_frames,
past_steps=cfg.past_frames, past_channels=6,
node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).to(device)
sd = {k: v for k, v in ckpt['interaction_graph_dict'].items()
if '_single_edge_index' not in k}
graph.load_state_dict(sd, strict=False)
graph.eval()
return model, init, graph
def make_beta_schedule(n=100, start=1e-4, end=5e-2):
return torch.linspace(start, end, n)
def extract(a, t, x):
out = torch.gather(a, 0, t.to(a.device))
return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1)))
@torch.no_grad()
def p_sample_accelerate(x, mask, cur_y, t, model, graph, use_v6_graph, sigma,
betas, alphas, alphas_bar_sqrt, one_minus_alphas_bar_sqrt):
t_tensor = torch.tensor([int(t)]).to(x.device)
eps_factor = ((1 - extract(alphas, t_tensor, cur_y))
/ extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y))
beta = extract(betas, t_tensor.repeat(x.shape[0]), cur_y)
eps_theta = model.generate_accelerate(cur_y, beta, x, mask)
if graph is not None:
abs_t = extract(alphas_bar_sqrt, t_tensor, cur_y)
am1_t = extract(one_minus_alphas_bar_sqrt, t_tensor, cur_y)
y0_hat = (cur_y - am1_t * eps_theta) / abs_t
delta = graph(y0_hat, x, int(t), sigma=sigma, A_override=x.size(0))
eps_theta = eps_theta - (abs_t / am1_t) * delta
mean = (1 / extract(alphas, t_tensor, cur_y).sqrt()) \
* (cur_y - eps_factor * eps_theta)
z = torch.randn_like(cur_y)
sigma_t = extract(betas, t_tensor, cur_y).sqrt()
return mean + sigma_t * z * 0.00001
@torch.no_grad()
def run(args):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cfg = Config(args.cfg, args.exp)
test_dset = SDDDataset(obs_len=cfg.past_frames,
pred_len=cfg.future_frames, split='test')
loader = DataLoader(test_dset, batch_size=1, shuffle=False,
num_workers=2, collate_fn=sdd_seq_collate)
exp_dir = cfg.model_dir # results_sdd/sdd/sdd/<info>/models
ckpt_path = cfg.model_path % args.epoch
print(f'Loading checkpoint: {ckpt_path}')
model, init, graph = build_models(
cfg, use_graph=args.use_graph, use_v6_graph=args.use_v6_graph,
ckpt_path=ckpt_path, device=device)
betas = make_beta_schedule().to(device)
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
abs_sqrt = torch.sqrt(alphas_prod)
one_minus_abs_sqrt = torch.sqrt(1 - alphas_prod)
traj_mean = torch.FloatTensor(cfg.traj_mean).to(device).view(1, 1, 1, 2)
traj_scale = float(cfg.traj_scale)
np.random.seed(0); random.seed(0)
torch.manual_seed(0); torch.cuda.manual_seed_all(0)
total_ade, total_fde, n = 0.0, 0.0, 0
T = cfg.future_frames
for data in loader:
pre = data['pre_motion_3D'].to(device) # [1, A, 8, 2]
fut = data['fut_motion_3D'].to(device)
A = pre.size(1)
initial_pos = pre[:, :, -1:]
past_abs = ((pre - traj_mean) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
past_rel = ((pre - initial_pos) / traj_scale).contiguous().view(-1, cfg.past_frames, 2)
past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1],
torch.zeros_like(past_rel[:, -1:])], dim=1)
past = torch.cat([past_abs, past_rel, past_vel], dim=-1)
fut_rel = ((fut - initial_pos) / traj_scale).contiguous().view(-1, T, 2)
mask = torch.ones(A, A).to(device)
sp, me, ve = init(past, mask)
ve = ve.clamp(min=-5, max=5)
sp = torch.exp(ve / 2)[..., None, None] * sp \
/ (sp.std(dim=1).mean(dim=(1, 2))[:, None, None, None] + 1e-6)
loc = sp + me[:, None]
sigma_in = ve if args.use_v6_graph else None
# leapfrog: 20 modes = 10+10 two halves, each 5 reverse steps
cur_y = loc[:, :10]
for i in reversed(range(NUM_Tau)):
cur_y = p_sample_accelerate(
past, mask, cur_y, i, model, graph, args.use_v6_graph, sigma_in,
betas, alphas, abs_sqrt, one_minus_abs_sqrt)
cur_y_ = loc[:, 10:]
for i in reversed(range(NUM_Tau)):
cur_y_ = p_sample_accelerate(
past, mask, cur_y_, i, model, graph, args.use_v6_graph, sigma_in,
betas, alphas, abs_sqrt, one_minus_abs_sqrt)
pred = torch.cat((cur_y_, cur_y), dim=1) # [A, 20, T, 2]
# target only (index 0) per scene, full-horizon ADE / final FDE, best_of_20
pred_0 = pred[0:1]
fut_0 = fut_rel[0:1]
dist = torch.norm(fut_0.unsqueeze(1) - pred_0, dim=-1) * traj_scale
ade = dist.mean(dim=-1).min(dim=-1)[0]
fde = dist[:, :, -1].min(dim=-1)[0]
total_ade += ade.sum().item()
total_fde += fde.sum().item()
n += 1
ade_px = total_ade / n * 50.0
fde_px = total_fde / n * 50.0
print(f'Epoch {args.epoch} MID-protocol: ADE={ade_px:.4f} FDE={fde_px:.4f} (n={n})')
if __name__ == '__main__':
p = argparse.ArgumentParser()
p.add_argument('--cfg', default='sdd/sdd')
p.add_argument('--exp', required=True, help='info tag e.g. baseline_v2')
p.add_argument('--epoch', type=int, required=True)
p.add_argument('--use_graph', action='store_true')
p.add_argument('--use_v6_graph', action='store_true')
args = p.parse_args()
run(args)
|