""" Stage-2 LED trainer for SDD. Variable-A scenes, batch_size=1, grad_accum. Optional graph module (--use_graph --use_v6_graph). Eval: only target agent (index 0) counts toward ADE/FDE. """ import os, time, torch, random, numpy as np import torch.nn as nn from utils.config import Config from utils.utils import print_log from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter 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 NUM_Tau = 5 class Trainer: def __init__(self, config): if torch.cuda.is_available(): torch.cuda.set_device(config.gpu) self.device = torch.device('cuda') if config.cuda else torch.device('cpu') self.cfg = Config(config.cfg, config.info) self.use_graph = bool(getattr(config, 'use_graph', False)) self.use_v6_graph = bool(getattr(config, 'use_v6_graph', False)) self.residual_on = getattr(config, 'residual_on', 'y0') self.grad_accum = getattr(config, 'grad_accum', 16) train_dset = SDDDataset(obs_len=self.cfg.past_frames, pred_len=self.cfg.future_frames, split='train') test_dset = SDDDataset(obs_len=self.cfg.past_frames, pred_len=self.cfg.future_frames, split='test') self.train_loader = DataLoader(train_dset, batch_size=1, shuffle=True, num_workers=2, collate_fn=sdd_seq_collate) self.test_loader = DataLoader(test_dset, batch_size=1, shuffle=False, num_workers=2, collate_fn=sdd_seq_collate) self.traj_mean = torch.FloatTensor(self.cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0) self.traj_scale = float(self.cfg.traj_scale) self.n_steps = self.cfg.diffusion.steps self.betas = self._make_beta_schedule( self.cfg.diffusion.beta_schedule, self.n_steps, self.cfg.diffusion.beta_start, self.cfg.diffusion.beta_end).cuda() self.alphas = 1 - self.betas self.alphas_prod = torch.cumprod(self.alphas, 0) self.alphas_bar_sqrt = torch.sqrt(self.alphas_prod) self.one_minus_alphas_bar_sqrt = torch.sqrt(1 - self.alphas_prod) self.model = CoreDenoisingModel(past_len=self.cfg.past_frames).cuda() ckpt_path = self.cfg.pretrained_core_denoising_model if not os.path.isfile(ckpt_path): raise FileNotFoundError(f'Missing pretrained denoiser: {ckpt_path}') self.model.load_state_dict(torch.load(ckpt_path, map_location='cpu')['model_dict']) self.model_initializer = InitializationModel( t_h=self.cfg.past_frames, d_h=6, t_f=self.cfg.future_frames, d_f=2, k_pred=20).cuda() params = list(self.model_initializer.parameters()) self.interaction_graph = None if self.use_graph: from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper self.interaction_graph = FutureInteractionGraphV6Wrapper( num_agents=64, future_steps=self.cfg.future_frames, past_steps=self.cfg.past_frames, past_channels=6, node_dim=128, top_n=5, num_denoise_steps=NUM_Tau).cuda() params += list(self.interaction_graph.parameters()) self.opt = torch.optim.AdamW(params, lr=config.learning_rate) self.scheduler = torch.optim.lr_scheduler.StepLR( self.opt, step_size=self.cfg.decay_step, gamma=self.cfg.decay_gamma) self.log = open(os.path.join(self.cfg.log_dir, 'log.txt'), 'a+') self.tb = SummaryWriter(log_dir=os.path.join(self.cfg.log_dir, 'tb')) self.global_step = 0 self._print_param(self.model, 'Core Denoiser') self._print_param(self.model_initializer, 'Initializer') if self.interaction_graph: self._print_param(self.interaction_graph, 'Graph') T = self.cfg.future_frames self.temporal_reweight = torch.FloatTensor( [(T + 1) - i for i in range(1, T + 1)]).cuda().unsqueeze(0).unsqueeze(0) / (T / 2) def _print_param(self, m, name): t = sum(p.numel() for p in m.parameters()) tr = sum(p.numel() for p in m.parameters() if p.requires_grad) print_log(f'[{name}] {tr}/{t}', self.log) def _make_beta_schedule(self, schedule, n, start, end): if schedule == 'linear': return torch.linspace(start, end, n) return torch.linspace(start, end, n) def _extract(self, a, t, x): out = torch.gather(a, 0, t.to(a.device)) return out.reshape(t.shape[0], *([1] * (len(x.shape) - 1))) def p_sample_accelerate(self, x, mask, cur_y, t, sigma=None): t_tensor = torch.tensor([int(t)]).cuda() eps_factor = ((1 - self._extract(self.alphas, t_tensor, cur_y)) / self._extract(self.one_minus_alphas_bar_sqrt, t_tensor, cur_y)) beta = self._extract(self.betas, t_tensor.repeat(x.shape[0]), cur_y) eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask) if self.interaction_graph is not None: abs_t = self._extract(self.alphas_bar_sqrt, t_tensor, cur_y) am1_t = self._extract(self.one_minus_alphas_bar_sqrt, t_tensor, cur_y) y0_hat = (cur_y - am1_t * eps_theta) / abs_t delta = self.interaction_graph( y0_hat, x, int(t), sigma=sigma, A_override=x.size(0)) eps_theta = eps_theta - (abs_t / am1_t) * delta mean = (1 / self._extract(self.alphas, t_tensor, cur_y).sqrt()) \ * (cur_y - eps_factor * eps_theta) z = torch.randn_like(cur_y) sigma_t = self._extract(self.betas, t_tensor, cur_y).sqrt() return mean + sigma_t * z * 0.00001 def p_sample_loop_accelerate(self, x, mask, loc, sigma=None): cur_y = loc[:, :10] for i in reversed(range(NUM_Tau)): cur_y = self.p_sample_accelerate(x, mask, cur_y, i, sigma=sigma) cur_y_ = loc[:, 10:] for i in reversed(range(NUM_Tau)): cur_y_ = self.p_sample_accelerate(x, mask, cur_y_, i, sigma=sigma) return torch.cat((cur_y_, cur_y), dim=1) def data_preprocess(self, data): pre = data['pre_motion_3D'].cuda() fut = data['fut_motion_3D'].cuda() A = pre.size(1) initial_pos = pre[:, :, -1:] past_abs = ((pre - self.traj_mean) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2) past_rel = ((pre - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.past_frames, 2) past_vel = torch.cat([past_rel[:, 1:] - past_rel[:, :-1], torch.zeros_like(past_rel[:, -1:])], dim=1) past_traj = torch.cat([past_abs, past_rel, past_vel], dim=-1) fut_traj = ((fut - initial_pos) / self.traj_scale).contiguous().view(-1, self.cfg.future_frames, 2) mask = torch.ones(A, A).cuda() return A, mask, past_traj, fut_traj def fit(self): for epoch in range(self.cfg.num_epochs): lt, ld, lu = self._train_epoch(epoch) print_log(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Epoch: {epoch}\t' f'Loss: {lt:.6f}\tDist: {ld:.6f}\tUnc: {lu:.6f}', self.log) self.tb.add_scalar('train/loss', lt, epoch) self.tb.add_scalar('train/loss_dist', ld, epoch) if (epoch + 1) % self.cfg.test_interval == 0: perf, n = self._test_epoch() # MID protocol: scale normalized ADE/FDE by 50 to report in pixels ade_px = perf['ADE'] / n * 50.0 fde_px = perf['FDE'] / n * 50.0 print_log(f'Epoch {epoch} Best Of 20: ADE: {ade_px:.4f} FDE: {fde_px:.4f}', self.log) self.tb.add_scalar('val/ADE_px', ade_px, epoch) self.tb.add_scalar('val/FDE_px', fde_px, epoch) cp = {'model_initializer_dict': self.model_initializer.state_dict()} if self.interaction_graph: cp['interaction_graph_dict'] = self.interaction_graph.state_dict() torch.save(cp, self.cfg.model_path % (epoch + 1)) self.scheduler.step() self.tb.flush(); self.tb.close() def _train_epoch(self, epoch): self.model.train(); self.model_initializer.train() if self.interaction_graph: self.interaction_graph.train() lt, ld, lu, cnt = 0, 0, 0, 0 self.opt.zero_grad() for i, data in enumerate(self.train_loader): A, mask, past, fut = self.data_preprocess(data) sp, me, ve = self.model_initializer(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 self.use_v6_graph else None gen = self.p_sample_loop_accelerate(past, mask, loc, sigma=sigma_in) loss_d = ((gen - fut.unsqueeze(1)).norm(p=2, dim=-1) * self.temporal_reweight).mean(dim=-1).min(dim=1)[0].mean() loss_u = (torch.exp(-ve) * (gen - fut.unsqueeze(1)).norm(p=2, dim=-1).mean(dim=(1, 2)) + ve).mean() loss = loss_d * 50 + loss_u (loss / self.grad_accum).backward() if (i + 1) % self.grad_accum == 0: params = list(self.model_initializer.parameters()) if self.interaction_graph: params += list(self.interaction_graph.parameters()) nn.utils.clip_grad_norm_(params, 1.0) self.opt.step(); self.opt.zero_grad() lt += loss.item(); ld += loss_d.item() * 50; lu += loss_u.item(); cnt += 1 self.global_step += 1 self.opt.step(); self.opt.zero_grad() return lt / cnt, ld / cnt, lu / cnt def _test_epoch(self): """MID-style SDD protocol: per-pedestrian full-horizon ADE (mean L2 over 12 future frames) and FDE (L2 at final frame), best_of_20 per pedestrian, then average across all evaluated pedestrians. Coordinates are already in MID's รท50 mean-centered space, so the final ADE/FDE is multiplied by 50 to report in pixels. Each scene in the SDD dataloader corresponds to one target pedestrian (index 0) + its neighbors; we evaluate only the target per scene so each pedestrian is counted exactly once (matches MID's get_timesteps_data qualification intent). """ T = self.cfg.future_frames perf = {'ADE': 0.0, 'FDE': 0.0} n = 0 np.random.seed(0); random.seed(0) torch.manual_seed(0); torch.cuda.manual_seed_all(0) self.model_initializer.eval() if self.interaction_graph: self.interaction_graph.eval() with torch.no_grad(): for data in self.test_loader: A, mask, past, fut = self.data_preprocess(data) sp, me, ve = self.model_initializer(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 self.use_v6_graph else None pred = self.p_sample_loop_accelerate(past, mask, loc, sigma=sigma_in) # MID protocol: only target agent (index 0) per scene pred_0 = pred[0:1] # [1, 20, T, 2] fut_0 = fut[0:1] # [1, T, 2] dist = torch.norm(fut_0.unsqueeze(1) - pred_0, dim=-1) * self.traj_scale # [1, 20, T] # best_of_20 per pedestrian, then full-horizon ADE / final FDE ade_per_ped = dist.mean(dim=-1).min(dim=-1)[0] # [1] fde_per_ped = dist[:, :, -1].min(dim=-1)[0] # [1] perf['ADE'] += ade_per_ped.sum().item() perf['FDE'] += fde_per_ped.sum().item() n += 1 return perf, n