""" LED trainer variant that augments the leapfrog denoising chain with a future-trajectory interaction graph (ported in spirit from MoFlow's FutureInteractionGraphV6 / MID graphv6_v5 ideas). At each of the NUM_Tau leapfrog reverse steps, the frozen core denoiser produces an epsilon prediction; we recover the implied y_0 estimate, run a small inter-agent graph on the predicted future trajectories, and add its output as a residual correction to epsilon. The graph is the only new trainable module besides the existing LED initializer. This file mirrors trainer/train_led_trajectory_augment_input.py so that the baseline script stays untouched and both variants can be run side-by-side. """ import os import time import torch import random import 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_nba import NBADataset, seq_collate from models.model_led_initializer import LEDInitializer as InitializationModel from models.model_diffusion import TransformerDenoisingModel as CoreDenoisingModel from models.future_interaction_graph import FutureInteractionGraph from models.future_interaction_graph_v6 import FutureInteractionGraphV6Wrapper 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) # ------------------------- prepare train/test data loader ------------------------- train_dset = NBADataset( obs_len=self.cfg.past_frames, pred_len=self.cfg.future_frames, training=True) self.train_loader = DataLoader( train_dset, batch_size=self.cfg.train_batch_size, shuffle=True, num_workers=4, collate_fn=seq_collate, pin_memory=True) test_dset = NBADataset( obs_len=self.cfg.past_frames, pred_len=self.cfg.future_frames, training=False) self.test_loader = DataLoader( test_dset, batch_size=self.cfg.test_batch_size, shuffle=False, num_workers=4, collate_fn=seq_collate, pin_memory=True) self.traj_mean = torch.FloatTensor(self.cfg.traj_mean).cuda().unsqueeze(0).unsqueeze(0).unsqueeze(0) self.traj_scale = self.cfg.traj_scale # ------------------------- define diffusion parameters ------------------------- self.n_steps = self.cfg.diffusion.steps self.betas = self.make_beta_schedule( schedule=self.cfg.diffusion.beta_schedule, n_timesteps=self.n_steps, start=self.cfg.diffusion.beta_start, end=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) # ------------------------- define models ------------------------- self.model = CoreDenoisingModel().cuda() model_cp = torch.load(self.cfg.pretrained_core_denoising_model, map_location='cpu') self.model.load_state_dict(model_cp['model_dict']) self.model_initializer = InitializationModel( t_h=10, d_h=6, t_f=20, d_f=2, k_pred=20).cuda() self.residual_on = getattr(config, 'residual_on', 'eps') self.use_sigma = bool(getattr(config, 'use_sigma', False)) self.use_v6_graph = bool(getattr(config, 'use_v6_graph', False)) self.uncertainty_weight = float(getattr(config, 'uncertainty_weight', 1.0)) top_n = getattr(config, 'top_n', 5) edge_mode = getattr(config, 'edge_mode', 'full') neighbor_mode = getattr(config, 'neighbor_mode', 'rag') if self.use_v6_graph: self.interaction_graph = FutureInteractionGraphV6Wrapper( num_agents=11, future_steps=self.cfg.future_frames, past_steps=self.cfg.past_frames, past_channels=6, node_dim=128, top_n=top_n, num_denoise_steps=NUM_Tau, edge_mode=edge_mode, neighbor_mode=neighbor_mode, ).cuda() else: self.interaction_graph = FutureInteractionGraph( num_agents=11, future_steps=self.cfg.future_frames, past_steps=self.cfg.past_frames, past_channels=6, node_dim=128, top_n=top_n, num_denoise_steps=NUM_Tau, use_sigma=self.use_sigma, ).cuda() self.opt = torch.optim.AdamW( list(self.model_initializer.parameters()) + list(self.interaction_graph.parameters()), lr=config.learning_rate, ) self.scheduler_model = torch.optim.lr_scheduler.StepLR( self.opt, step_size=self.cfg.decay_step, gamma=self.cfg.decay_gamma) self.resume_epoch = int(getattr(config, 'resume_epoch', 0)) if self.resume_epoch > 0: cp_path = self.cfg.model_path % self.resume_epoch cp = torch.load(cp_path, map_location='cpu') self.model_initializer.load_state_dict(cp['model_initializer_dict']) self.interaction_graph.load_state_dict(cp['interaction_graph_dict']) for _ in range(self.resume_epoch): self.scheduler_model.step() # ------------------------- prepare logs ------------------------- 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_model_param(self.model, name='Core Denoising Model') self.print_model_param(self.model_initializer, name='Initialization Model') self.print_model_param(self.interaction_graph, name='Future Interaction Graph') self.temporal_reweight = torch.FloatTensor( [21 - i for i in range(1, 21)]).cuda().unsqueeze(0).unsqueeze(0) / 10 def print_model_param(self, model: nn.Module, name: str = 'Model') -> None: total_num = sum(p.numel() for p in model.parameters()) trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad) print_log("[{}] Trainable/Total: {}/{}".format(name, trainable_num, total_num), self.log) def make_beta_schedule(self, schedule: str = 'linear', n_timesteps: int = 1000, start: float = 1e-5, end: float = 1e-2) -> torch.Tensor: if schedule == 'linear': betas = torch.linspace(start, end, n_timesteps) elif schedule == "quad": betas = torch.linspace(start ** 0.5, end ** 0.5, n_timesteps) ** 2 elif schedule == "sigmoid": betas = torch.linspace(-6, 6, n_timesteps) betas = torch.sigmoid(betas) * (end - start) + start return betas def extract(self, input, t, x): shape = x.shape out = torch.gather(input, 0, t.to(input.device)) reshape = [t.shape[0]] + [1] * (len(shape) - 1) return out.reshape(*reshape) # ------------------------------------------------------------------ # Leapfrog reverse step with graph residual on epsilon # ------------------------------------------------------------------ def p_sample_accelerate(self, x, mask, cur_y, t, sigma=None): step_idx = int(t) t = torch.tensor([t]).cuda() eps_factor = ((1 - self.extract(self.alphas, t, cur_y)) / self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y)) beta = self.extract(self.betas, t.repeat(x.shape[0]), cur_y) eps_theta = self.model.generate_accelerate(cur_y, beta, x, mask) # Implied y_0 estimate for the graph input. alpha_bar_sqrt_t = self.extract(self.alphas_bar_sqrt, t, cur_y) one_minus_abs_t = self.extract(self.one_minus_alphas_bar_sqrt, t, cur_y) y0_hat = (cur_y - one_minus_abs_t * eps_theta) / alpha_bar_sqrt_t delta = self.interaction_graph(y0_hat, x, step_idx, sigma=sigma) if self.residual_on == 'eps': eps_theta = eps_theta + delta else: eps_theta = eps_theta - (alpha_bar_sqrt_t / one_minus_abs_t) * delta mean = (1 / self.extract(self.alphas, t, cur_y).sqrt()) * (cur_y - (eps_factor * eps_theta)) z = torch.randn_like(cur_y).to(x.device) sigma_t = self.extract(self.betas, t, cur_y).sqrt() sample = mean + sigma_t * z * 0.00001 return sample 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) prediction_total = torch.cat((cur_y_, cur_y), dim=1) return prediction_total # ------------------------------------------------------------------ # Training / evaluation # ------------------------------------------------------------------ def fit(self): for epoch in range(self.resume_epoch, self.cfg.num_epochs): loss_total, loss_distance, loss_uncertainty = self._train_single_epoch(epoch) print_log('[{}] Epoch: {}\t\tLoss: {:.6f}\tLoss Dist.: {:.6f}\tLoss Uncertainty: {:.6f}'.format( time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), epoch, loss_total, loss_distance, loss_uncertainty), self.log) self.tb.add_scalar('train_epoch/loss_total', loss_total, epoch) self.tb.add_scalar('train_epoch/loss_dist_x50', loss_distance, epoch) self.tb.add_scalar('train_epoch/loss_uncertainty', loss_uncertainty, epoch) self.tb.add_scalar('train_epoch/lr', self.opt.param_groups[0]['lr'], epoch) if (epoch + 1) % self.cfg.test_interval == 0: performance, samples = self._test_single_epoch() for time_i in range(4): ade = performance['ADE'][time_i] / samples fde = performance['FDE'][time_i] / samples print_log('--ADE({}s): {:.4f}\t--FDE({}s): {:.4f}'.format( time_i + 1, ade, time_i + 1, fde), self.log) self.tb.add_scalar('val/ADE_{}s'.format(time_i + 1), ade, epoch) self.tb.add_scalar('val/FDE_{}s'.format(time_i + 1), fde, epoch) cp_path = self.cfg.model_path % (epoch + 1) model_cp = { 'model_initializer_dict': self.model_initializer.state_dict(), 'interaction_graph_dict': self.interaction_graph.state_dict(), } torch.save(model_cp, cp_path) self.scheduler_model.step() self.tb.flush() self.tb.close() def data_preprocess(self, data): batch_size = data['pre_motion_3D'].shape[0] traj_mask = torch.zeros(batch_size * 11, batch_size * 11).cuda() for i in range(batch_size): traj_mask[i * 11:(i + 1) * 11, i * 11:(i + 1) * 11] = 1. initial_pos = data['pre_motion_3D'].cuda()[:, :, -1:] past_traj_abs = ((data['pre_motion_3D'].cuda() - self.traj_mean) / self.traj_scale).contiguous().view(-1, 10, 2) past_traj_rel = ((data['pre_motion_3D'].cuda() - initial_pos) / self.traj_scale).contiguous().view(-1, 10, 2) past_traj_vel = torch.cat( (past_traj_rel[:, 1:] - past_traj_rel[:, :-1], torch.zeros_like(past_traj_rel[:, -1:])), dim=1) past_traj = torch.cat((past_traj_abs, past_traj_rel, past_traj_vel), dim=-1) fut_traj = ((data['fut_motion_3D'].cuda() - initial_pos) / self.traj_scale).contiguous().view(-1, 20, 2) return batch_size, traj_mask, past_traj, fut_traj def _train_single_epoch(self, epoch): self.model.train() self.model_initializer.train() self.interaction_graph.train() loss_total, loss_dt, loss_dc, count = 0, 0, 0, 0 for data in self.train_loader: batch_size, traj_mask, past_traj, fut_traj = self.data_preprocess(data) sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask) # --- σ 안정화 가드 (기본 전부 off => 원본과 비트 단위로 동일) -------------- # 배경: --use_sigma 로 logvar 를 그래프에 넣으면 NLL 이외의 gradient 경로가 # 하나 더 생긴다. logvar 가 음수로 밀리면 exp(-logvar) 가 폭주하고 # exp(logvar/2)*x / std(x) 의 분모도 0 으로 가서 NaN 이 된다. # (실제로 LED full SRA σ ON 이 epoch 4 에서 이렇게 죽었다) _clamp = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0) if _clamp > 0: variance_estimation = variance_estimation.clamp(-_clamp, _clamp) sample_prediction = (torch.exp(variance_estimation / 2)[..., None, None] * sample_prediction / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None] .clamp_min(float(os.environ.get('LED_STD_EPS', 0.0) or 0.0))) loc = sample_prediction + mean_estimation[:, None] sigma_input = variance_estimation if self.use_sigma else None # σ 를 '게이트 신호'로만 쓰고 분산 헤드로 역전파하지 않는 옵션 if sigma_input is not None and os.environ.get('LED_SIGMA_DETACH', '') not in ('', '0', 'false', 'False'): sigma_input = sigma_input.detach() generated_y = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input) loss_dist = ((generated_y - fut_traj.unsqueeze(dim=1)).norm(p=2, dim=-1) * self.temporal_reweight).mean(dim=-1).min(dim=1)[0].mean() loss_uncertainty = (torch.exp(-variance_estimation) * (generated_y - fut_traj.unsqueeze(dim=1)).norm(p=2, dim=-1).mean(dim=(1, 2)) + variance_estimation).mean() # logvar 추이를 남긴다 (NaN 이 나면 원인을 사후에 알 수 있게) if count % 200 == 0: with torch.no_grad(): _v = variance_estimation.detach() self.tb.add_scalar('sigma/logvar_min', float(_v.min()), self.global_step) self.tb.add_scalar('sigma/logvar_max', float(_v.max()), self.global_step) self.tb.add_scalar('sigma/logvar_mean', float(_v.mean()), self.global_step) self.tb.add_scalar('sigma/pred_std_min', float(sample_prediction.std(dim=1).mean(dim=(1, 2)).min()), self.global_step) loss = loss_dist * 50 + self.uncertainty_weight * loss_uncertainty loss_total += loss.item() loss_dt += loss_dist.item() * 50 loss_dc += loss_uncertainty.item() self.opt.zero_grad() loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_( list(self.model_initializer.parameters()) + list(self.interaction_graph.parameters()), 1., ) self.opt.step() self.tb.add_scalar('train_step/loss_total', loss.item(), self.global_step) self.tb.add_scalar('train_step/loss_dist_x50', loss_dist.item() * 50, self.global_step) self.tb.add_scalar('train_step/loss_uncertainty', loss_uncertainty.item(), self.global_step) self.tb.add_scalar('train_step/grad_norm', float(grad_norm), self.global_step) self.global_step += 1 count += 1 if self.cfg.debug and count == 2: break return loss_total / count, loss_dt / count, loss_dc / count def _test_single_epoch(self): performance = {'FDE': [0, 0, 0, 0], 'ADE': [0, 0, 0, 0]} samples = 0 def prepare_seed(rand_seed): np.random.seed(rand_seed) random.seed(rand_seed) torch.manual_seed(rand_seed) torch.cuda.manual_seed_all(rand_seed) prepare_seed(0) self.model_initializer.eval() self.interaction_graph.eval() with torch.no_grad(): for data in self.test_loader: batch_size, traj_mask, past_traj, fut_traj = self.data_preprocess(data) sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask) _c = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0) if _c > 0: variance_estimation = variance_estimation.clamp(-_c, _c) sample_prediction = (torch.exp(variance_estimation / 2)[..., None, None] * sample_prediction / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None]) loc = sample_prediction + mean_estimation[:, None] sigma_input = variance_estimation if self.use_sigma else None pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input) fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1) distances = torch.norm(fut_traj - pred_traj, dim=-1) * self.traj_scale for time_i in range(1, 5): ade = (distances[:, :, :5 * time_i]).mean(dim=-1).min(dim=-1)[0].sum() fde = (distances[:, :, 5 * time_i - 1]).min(dim=-1)[0].sum() performance['ADE'][time_i - 1] += ade.item() performance['FDE'][time_i - 1] += fde.item() samples += distances.shape[0] return performance, samples def test_single_model(self): model_path = './results/checkpoints/led_graph.p' ckpt = torch.load(model_path, map_location=torch.device('cpu')) self.model_initializer.load_state_dict(ckpt['model_initializer_dict']) if 'interaction_graph_dict' in ckpt: self.interaction_graph.load_state_dict(ckpt['interaction_graph_dict']) else: print_log('WARNING: checkpoint has no interaction_graph_dict; ' 'using zero-initialized graph (equivalent to baseline LED).', log=self.log) performance = {'FDE': [0, 0, 0, 0], 'ADE': [0, 0, 0, 0]} samples = 0 print_log(model_path, log=self.log) def prepare_seed(rand_seed): np.random.seed(rand_seed) random.seed(rand_seed) torch.manual_seed(rand_seed) torch.cuda.manual_seed_all(rand_seed) prepare_seed(0) self.model_initializer.eval() self.interaction_graph.eval() with torch.no_grad(): for data in self.test_loader: batch_size, traj_mask, past_traj, fut_traj = self.data_preprocess(data) sample_prediction, mean_estimation, variance_estimation = self.model_initializer(past_traj, traj_mask) _c = float(os.environ.get('LED_SIGMA_CLAMP', 0.0) or 0.0) if _c > 0: variance_estimation = variance_estimation.clamp(-_c, _c) sample_prediction = (torch.exp(variance_estimation / 2)[..., None, None] * sample_prediction / sample_prediction.std(dim=1).mean(dim=(1, 2))[:, None, None, None]) loc = sample_prediction + mean_estimation[:, None] sigma_input = variance_estimation if self.use_sigma else None pred_traj = self.p_sample_loop_accelerate(past_traj, traj_mask, loc, sigma=sigma_input) fut_traj = fut_traj.unsqueeze(1).repeat(1, 20, 1, 1) distances = torch.norm(fut_traj - pred_traj, dim=-1) * self.traj_scale for time_i in range(1, 5): ade = (distances[:, :, :5 * time_i]).mean(dim=-1).min(dim=-1)[0].sum() fde = (distances[:, :, 5 * time_i - 1]).min(dim=-1)[0].sum() performance['ADE'][time_i - 1] += ade.item() performance['FDE'][time_i - 1] += fde.item() samples += distances.shape[0] for time_i in range(4): print_log('--ADE({}s): {:.4f}\t--FDE({}s): {:.4f}'.format( time_i + 1, performance['ADE'][time_i] / samples, time_i + 1, performance['FDE'][time_i] / samples), log=self.log)