| """ |
| Stage-1 pretraining for LED on SDD. |
| Variable-A scenes with batch_size=1 + gradient accumulation. |
| past_frames=8, future_frames=12. Pixel coordinates. |
| """ |
|
|
| import os, time, torch, 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_diffusion import TransformerDenoisingModel as CoreDenoisingModel |
|
|
|
|
| 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.grad_accum = getattr(config, 'grad_accum', 32) |
|
|
| 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.val_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() |
|
|
| pre_lr = float(self.cfg.pretrain['lr']) |
| self.pre_epochs = int(self.cfg.pretrain['num_epochs']) |
| self.opt = torch.optim.AdamW(self.model.parameters(), lr=pre_lr) |
| self.scheduler = torch.optim.lr_scheduler.StepLR( |
| self.opt, step_size=int(self.cfg.pretrain.get('decay_step', 30)), |
| gamma=float(self.cfg.pretrain.get('decay_gamma', 0.5))) |
|
|
| 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.ckpt_path = self.cfg.pretrained_core_denoising_model |
| total = sum(p.numel() for p in self.model.parameters()) |
| print_log(f'Core Denoiser params: {total:,}', self.log) |
|
|
| def _make_beta_schedule(self, schedule, n, start, end): |
| if schedule == 'linear': return torch.linspace(start, end, n) |
| elif schedule == 'quad': return torch.linspace(start**0.5, end**0.5, n)**2 |
| 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 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 noise_estimation_loss(self, x, y_0, mask): |
| B = x.shape[0] |
| t = torch.randint(0, self.n_steps, size=(B // 2 + 1,)).to(x.device) |
| t = torch.cat([t, self.n_steps - t - 1], dim=0)[:B] |
| a = self._extract(self.alphas_bar_sqrt, t, y_0) |
| beta = self._extract(self.betas, t, y_0) |
| am1 = self._extract(self.one_minus_alphas_bar_sqrt, t, y_0) |
| e = torch.randn_like(y_0) |
| y = y_0 * a + e * am1 |
| out = self.model(y, beta, x, mask) |
| return (e - out).square().mean() |
|
|
| def fit(self): |
| best_val = float('inf') |
| for epoch in range(self.pre_epochs): |
| self.model.train() |
| loss_sum, n = 0.0, 0 |
| self.opt.zero_grad() |
| for i, data in enumerate(self.train_loader): |
| _, mask, past, fut = self.data_preprocess(data) |
| loss = self.noise_estimation_loss(past, fut, mask) |
| (loss / self.grad_accum).backward() |
| if (i + 1) % self.grad_accum == 0: |
| nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) |
| self.opt.step(); self.opt.zero_grad() |
| loss_sum += loss.item(); n += 1 |
| self.tb.add_scalar('pretrain_step/loss', loss.item(), self.global_step) |
| self.global_step += 1 |
| self.opt.step(); self.opt.zero_grad() |
|
|
| train_loss = loss_sum / max(1, n) |
| print_log(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Pretrain Epoch {epoch} train_mse={train_loss:.6f}', self.log) |
| self.tb.add_scalar('pretrain_epoch/train_loss', train_loss, epoch) |
|
|
| if (epoch + 1) % 5 == 0: |
| self.model.eval() |
| val_sum, val_n = 0.0, 0 |
| with torch.no_grad(): |
| for data in self.val_loader: |
| _, mask, past, fut = self.data_preprocess(data) |
| val_sum += self.noise_estimation_loss(past, fut, mask).item() |
| val_n += 1 |
| val_loss = val_sum / max(1, val_n) |
| print_log(f' val_mse={val_loss:.6f}', self.log) |
| self.tb.add_scalar('pretrain_epoch/val_loss', val_loss, epoch) |
| if val_loss < best_val: |
| best_val = val_loss |
| os.makedirs(os.path.dirname(self.ckpt_path), exist_ok=True) |
| torch.save({'model_dict': self.model.state_dict(), 'epoch': epoch}, self.ckpt_path) |
| print_log(f' -> saved {self.ckpt_path}', self.log) |
|
|
| self.scheduler.step() |
| self.tb.flush(); self.tb.close() |
|
|
