import copy import functools import os import time from types import SimpleNamespace import numpy as np import blobfile as bf import torch from torch.optim import AdamW from diffusion import logger from utils import dist_util from diffusion.fp16_util import MixedPrecisionTrainer from diffusion.resample import LossAwareSampler, UniformSampler from tqdm import tqdm from diffusion.resample import create_named_schedule_sampler # from data_loaders.humanml.networks.evaluator_wrapper import EvaluatorMDMWrapper # from eval import eval_humanml, eval_humanact12_uestc # from data_loaders.get_data import get_dataset_loader # For ImageNet experiments, this was a good default value. # We found that the lg_loss_scale quickly climbed to # 20-21 within the first ~1K steps of training. ## intial log loss scale ## INITIAL_LOG_LOSS_SCALE = 20.0 # do not # class TrainLoop: def __init__(self, args, train_platform, model, diffusion, data): self.args = args self.dataset = args.dataset self.train_platform = train_platform self.model = model self.diffusion = diffusion if self.args.nprocs > 1: self.cond_mode = model.module.cond_mode else: self.cond_mode = model.cond_mode self.data = data self.batch_size = args.batch_size self.microbatch = args.batch_size # deprecating this option self.lr = args.lr self.log_interval = args.log_interval self.save_interval = args.save_interval self.resume_checkpoint = args.resume_checkpoint self.use_fp16 = False # deprecating this option self.fp16_scale_growth = 1e-3 # deprecating this option self.weight_decay = args.weight_decay self.lr_anneal_steps = args.lr_anneal_steps self.step = 0 # self.resume_step = 0 self.resume_step = False self.global_batch = self.batch_size # * dist.get_world_size() self.num_steps = args.num_steps self.num_epochs = self.num_steps // len(self.data) + 1 self.sync_cuda = False # torch.cuda.is_available() if self.args.finetune_with_cond: # finetune_with_cond -> self._load_and_sync_parameters_cond() # load parameters here print(f"Setting trans linear layer to zero for conditioning...") self.model.set_trans_linear_layer_to_zero() # else: # finetune_with_cond self._load_and_sync_parameters() self.mp_trainer = MixedPrecisionTrainer( # mixed model=self.model, # use_fp16=self.use_fp16, fp16_scale_growth=self.fp16_scale_growth, args=args, ) self.save_dir = args.save_dir self.overwrite = args.overwrite self.opt = AdamW( self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay ) if self.resume_step and not args.not_load_opt: self._load_optimizer_state() # Model was resumed, either due to a restart or a checkpoint # being specified at the command line. print(f"dist_utils: {dist_util.dev()}") self.device = torch.device("cpu") # if torch.cuda.is_available() and dist_util.dev() != 'cpu': # self.device = torch.device(dist_util.dev()) self.schedule_sampler_type = 'uniform' self.schedule_sampler = create_named_schedule_sampler(self.schedule_sampler_type, diffusion) self.eval_wrapper, self.eval_data, self.eval_gt_data = None, None, None # if args.dataset in ['kit', 'humanml', 'motion_ours'] and args.eval_during_training: # mm_num_samples = 0 # mm is super slow hence we won't run it during training # mm_num_repeats = 0 # mm is super slow hence we won't run it during training # gen_loader = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None, # split=args.eval_split, # hml_mode='eval') # self.eval_gt_data = get_dataset_loader(name=args.dataset, batch_size=args.eval_batch_size, num_frames=None, # split=args.eval_split, # hml_mode='gt') # self.eval_wrapper = EvaluatorMDMWrapper(args.dataset, dist_util.dev()) # self.eval_data = { # 'test': lambda: eval_humanml.get_mdm_loader( # model, diffusion, args.eval_batch_size, # gen_loader, mm_num_samples, mm_num_repeats, gen_loader.dataset.opt.max_motion_length, # args.eval_num_samples, scale=1., # ) # } self.use_ddp = False if self.args.nprocs == 1 else True self.ddp_model = self.model def safe_load_ckpt(self, model, state_dicts): ori_dict = state_dicts part_dict = dict() model_dict = model.state_dict() tot_params_n = 0 for k in ori_dict: if self.args.resume_diff: if k in model_dict: v = ori_dict[k] part_dict[k] = v tot_params_n += 1 else: if k in model_dict and "denoising" not in k: v = ori_dict[k] part_dict[k] = v tot_params_n += 1 model_dict.update(part_dict) model.load_state_dict(model_dict) print(f"Resume glb-backbone finished!! Total number of parameters: {tot_params_n}.") # def _load_and_sync_parameters_cond(self): resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint if resume_checkpoint: self.resume_step = parse_resume_step_from_filename(resume_checkpoint) logger.log(f"loading model from checkpoint: {resume_checkpoint}...") state_dicts = dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev() ) if self.args.diff_basejtsrel: # if self.args.finetune_with_cond_rel: model_dict = self.model.state_dict() # elif self.args.finetune_with_cond_jtsobj: model_dict.update(state_dicts) self.model.load_state_dict(model_dict) if self.args.finetune_with_cond_jtsobj: # finetune_with_cond_jtsobj --> finetune_with_cond_jtsobj # cond_joints_offset_input_process <- joints_offset_input_process; cond_sequence_pos_encoder <- sequence_pos_encoder; cond_seqTransEncoder <- seqTransEncoder self.model.cond_joints_offset_input_process.load_state_dict(self.model.joints_offset_input_process.state_dict()) self.model.cond_sequence_pos_encoder.load_state_dict(self.model.sequence_pos_encoder.state_dict()) self.model.cond_seqTransEncoder.load_state_dict(self.model.seqTransEncoder.state_dict()) else: raise ValueError(f"Must have diff_basejtsrel setting, others not implemented yet!") # self.safe_load_ckpt(self.model, # dist_util.load_state_dict( # resume_checkpoint, map_location=dist_util.dev() # ) # ) def _load_and_sync_parameters(self): resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint if resume_checkpoint: self.resume_step = parse_resume_step_from_filename(resume_checkpoint) logger.log(f"loading model from checkpoint: {resume_checkpoint}...") # self.model.load_state_dict( # dist_util.load_state_dict( # resume_checkpoint, map_location=dist_util.dev() # ) # ) self.safe_load_ckpt(self.model, dist_util.load_state_dict( resume_checkpoint, map_location=dist_util.dev() ) ) def _load_optimizer_state(self): main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint opt_checkpoint = bf.join( bf.dirname(main_checkpoint), f"opt{self.resume_step:09}.pt" ) if bf.exists(opt_checkpoint): logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") state_dict = dist_util.load_state_dict( opt_checkpoint, map_location=dist_util.dev() ) self.opt.load_state_dict(state_dict) def run_loop(self): for epoch in range(self.num_epochs): print(f'Starting epoch {epoch}') for batch in tqdm(self.data): # for motion, cond in tqdm(self.data): ## motion; cond; data ## if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps): break # print(f"motion: {motion.size()}, ") ## motion.to(self.device) # motion = motion.to(self.device) # cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()} # batch = { # key: val.to(self.device) if torch.is_tensor(val) else ([subval.to(self.device) for subval in val] if isinstance(val, list) else val) for key, val in batch.items() # } for k in batch: if isinstance(batch[k], torch.Tensor): batch[k] = batch[k].to(self.device) elif isinstance(batch[k], list): batch[k] = [subval.to(self.device) if isinstance(subval, torch.Tensor) else subval for subval in batch[k]] else: batch[k] = batch[k] ## run current motion and cond ## ## run step ## self.run_step(batch) ## run step for the motion and cond ## ## ===== log useful things ==== ## if self.step % self.log_interval == 0: # loss_dict = logger.get_current().name2val print('step[{}]: loss[{:0.5f}]'.format(self.step+self.resume_step, loss_dict["loss"])) for k in loss_dict: v = loss_dict[k] if k in ['rel_pred_loss', 'dist_pred_loss', 'dec_e_along_normals_loss', 'dec_e_vt_normals_loss', 'joints_pred_loss', 'jts_pred_loss', 'jts_latent_denoising_loss', 'basejtsrel_pred_loss', 'basejtsrel_latent_denoising_loss', 'basejtse_along_normals_pred_loss', 'basejtse_vt_normals_pred_loss', 'basejtse_latent_denoising_loss', "KL_loss", "avg_joints_pred_loss", "basejtrel_denoising_loss", "avgjts_denoising_loss"]: ## avg_joints_pred_loss # avg joints pred loss # print(f"\t{k}: {loss_dict[k].mean().item() if isinstance(loss_dict[k], torch.Tensor) else loss_dict[k]}") if k in ['step', 'samples'] or '_q' in k: # step samples # continue else: self.train_platform.report_scalar(name=k, value=v, iteration=self.step, group_name='Loss') # for k,v in logger.get_current().name2val.items(): # if k == 'loss': # print('step[{}]: loss[{:0.5f}]'.format(self.step+self.resume_step, v)) # if k in ['step', 'samples'] or '_q' in k: # step samples # # continue # else: # self.train_platform.report_scalar(name=k, value=v, iteration=self.step, group_name='Loss') ## ===== save checkpoints ===== ## if self.step % self.save_interval == 0: ## save; model.eval; self.save() if self.args.nprocs > 1: self.model.module.eval() else: self.model.eval() self.evaluate() if self.args.nprocs > 1: self.model.module.train() else: self.model.train() # Run for a finite amount of time in integration tests. if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: return self.step += 1 if not (not self.lr_anneal_steps or self.step + self.resume_step < self.lr_anneal_steps): break # Save the last checkpoint if it wasn't already saved. if (self.step - 1) % self.save_interval != 0: self.save() self.evaluate() def evaluate(self): if not self.args.eval_during_training: return start_eval = time.time() if self.eval_wrapper is not None: print('Running evaluation loop: [Should take about 90 min]') log_file = os.path.join(self.save_dir, f'eval_humanml_{(self.step + self.resume_step):09d}.log') diversity_times = 300 mm_num_times = 0 # mm is super slow hence we won't run it during training eval_dict = eval_humanml.evaluation( self.eval_wrapper, self.eval_gt_data, self.eval_data, log_file, replication_times=self.args.eval_rep_times, diversity_times=diversity_times, mm_num_times=mm_num_times, run_mm=False) print(eval_dict) for k, v in eval_dict.items(): if k.startswith('R_precision'): for i in range(len(v)): self.train_platform.report_scalar(name=f'top{i + 1}_' + k, value=v[i], iteration=self.step + self.resume_step, group_name='Eval') else: self.train_platform.report_scalar(name=k, value=v, iteration=self.step + self.resume_step, group_name='Eval') elif self.dataset in ['humanact12', 'uestc']: eval_args = SimpleNamespace(num_seeds=self.args.eval_rep_times, num_samples=self.args.eval_num_samples, batch_size=self.args.eval_batch_size, device=self.device, guidance_param = 1, dataset=self.dataset, unconstrained=self.args.unconstrained, model_path=os.path.join(self.save_dir, self.ckpt_file_name())) eval_dict = eval_humanact12_uestc.evaluate(eval_args, model=self.model, diffusion=self.diffusion, data=self.data.dataset) print(f'Evaluation results on {self.dataset}: {sorted(eval_dict["feats"].items())}') for k, v in eval_dict["feats"].items(): if 'unconstrained' not in k: self.train_platform.report_scalar(name=k, value=np.array(v).astype(float).mean(), iteration=self.step, group_name='Eval') else: self.train_platform.report_scalar(name=k, value=np.array(v).astype(float).mean(), iteration=self.step, group_name='Eval Unconstrained') end_eval = time.time() print(f'Evaluation time: {round(end_eval-start_eval)/60}min') def run_step(self, batch): self.forward_backward(batch) ## forward self.mp_trainer.optimize(self.opt) self._anneal_lr() self.log_step() def forward_backward(self, batch): self.mp_trainer.zero_grad() for i in range(0, batch['base_pts'].shape[0], self.microbatch): # print(f"batch_device: {batch['base_pts'].device}") ## base pts device # Eliminates the microbatch feature assert i == 0 assert self.microbatch == self.batch_size micro = batch # micro_cond = cond ## micro-batch # base_pts; base_pts # last_batch = (i + self.microbatch) >= batch['base_pts'].shape[0] t, weights = self.schedule_sampler.sample(micro['base_pts'].shape[0], dist_util.dev()) # print(f"t: {t.size()}, weights: {weights.size()}, t_device: {t.device}, weights_device: {weights.device}") # compute_losses = functools.partial( # self.diffusion.training_losses, # self.ddp_model, # micro, # [bs, ch, image_size, image_size] # t, # [bs](int) sampled timesteps # model_kwargs={'y': batch}, # # dataset=self.data.dataset # ) # # if last_batch or not self.use_ddp: # # losses = compute_losses() ## compute lossses # # else: # # with self.ddp_model.no_sync(): # # losses = compute_losses() # if not self.use_ddp: # losses = compute_losses() ## compute lossses # else: # with self.ddp_model.no_sync(): # losses = compute_losses() ### training losses ### losses = self.diffusion.training_losses( self.ddp_model, micro, # [bs, ch, image_size, image_size] t, # [bs](int) sampled timesteps model_kwargs={'y': batch}, dataset=self.data.dataset ) # loss aware sampler # if isinstance(self.schedule_sampler, LossAwareSampler): self.schedule_sampler.update_with_local_losses( t, losses["loss"].detach() ) # print(losses["loss"].size(), f"weights: {weights.size()}") loss = (losses["loss"] * weights).mean() if self.args.nprocs > 1: torch.distributed.barrier() dist_util.reduce_mean(loss, self.args.nprocs) ## args nprocs ## log_loss_dict( self.diffusion, t, {k: v * weights for k, v in losses.items()} ) self.mp_trainer.backward(loss) def predict_single_step(self, batch, use_t=None): # self.mp_trainer.zero_grad() # use_t is not Noen tot_samples = [] tot_targets = [] tot_dec_disp_e_along_normals = [] tot_dec_disp_e_vt_normals = [] tot_pred_joints_quant = [] # for i in range(0, batch['base_pts'].shape[0], self.microbatch): # Eliminates the microbatch feature assert i == 0 assert self.microbatch == self.batch_size micro = batch # ## micro batch ## rhand_joints = micro['rhand_joints'] # micro_cond = cond # micro_cond and cond ## ## micro-batch ## last_batch = (i + self.microbatch) >= batch['base_pts'].shape[0] t, weights = self.schedule_sampler.sample(micro['base_pts'].shape[0], dist_util.dev()) if use_t is not None: t = torch.zeros_like(t) + use_t # batch: bsz x nnjoints x 3 x nnframes # ## === original sampling === ## # terms, model_output, target, t = self.diffusion.predict_sample_single_step(self.ddp_model, micro, t, model_kwargs=micro_cond, noise=None, dataset=self.data.dataset) ## restricted by those things ## ### use p_sample_loop from the diffusion model ### sample_fn = self.diffusion.p_sample_loop samples = sample_fn( self.ddp_model, rhand_joints.shape, clip_denoised=False, model_kwargs=micro, skip_timesteps=0, init_image=micro, progress=True, dump_steps=None, noise=None, ## noise ## # # const_noise=False, # whether to cond on noise ## const_noise=self.args.const_noise, ## const noise ! st_timestep=use_t, ) # sample either as joints or as relative positions for each base pts # tot_samples.append(samples['sampled_rhand_joints']) # tot_samples = tot_samples + samples # samples rhand_joints; targets rhand_joints ### add rhand joints tot_targets.append(micro['rhand_joints']) if 'e_disp_rel_to_base_along_normals' in samples: tot_dec_disp_e_along_normals.append(samples['e_disp_rel_to_base_along_normals']) tot_dec_disp_e_vt_normals.append(samples['e_disp_rel_to_baes_vt_normals']) if 'pred_joint_quants' in samples: tot_pred_joints_quant.append(samples['pred_joint_quants']) # tot_targets.append(samples['rhand_joints']) # all of them target at joints samples ## model_output = torch.cat(tot_samples, dim=0) # model_output = tot_samples target = torch.cat(tot_targets, dim=0) if len(tot_dec_disp_e_along_normals) > 0: tot_dec_disp_e_along_normals = torch.cat(tot_dec_disp_e_along_normals, dim=0) tot_dec_disp_e_vt_normals = torch.cat(tot_dec_disp_e_vt_normals, dim=0) ### tot_dec_disp_e_vt_normals # if len(tot_pred_joints_quant) > 0: tot_pred_joints_quant = torch.cat(tot_pred_joints_quant, dim=0) # print(f"Returning with model_output; {model_output.size()}, target: {target.size()}") print(f"Returning with target: {target.size()}") ### returning the samples and tarets ### if isinstance(tot_pred_joints_quant, torch.Tensor): return model_output, target, tot_pred_joints_quant elif isinstance(tot_dec_disp_e_along_normals, torch.Tensor): return model_output, target, tot_dec_disp_e_along_normals, tot_dec_disp_e_vt_normals else: return model_output, target # return model_output, target ### predict from data ### def predict_from_data(self): # for epoch in range(self.num_epochs): # # print(f'Starting epoch {epoch}') # the ## ==== a single pass for a single sequence ==== ## tot_model_outputs = [] tot_targets = [] tot_st_idxes = [] tot_ed_idxes = [] tot_pert_verts = [] tot_verts = [] tot_dec_disp_e_along_normals = [] tot_dec_disp_e_vt_normals = [] ## motion; cond; data ## tot_pred_joints_quant = [] for batch in tqdm(self.data): # batch data # for k in batch: if isinstance(batch[k], torch.Tensor): batch[k] = batch[k].to(self.device) elif isinstance(batch[k], list): # batch[k] = [subval.to(self.device) for subval in batch[k]] batch[k] = [subval.to(self.device) if isinstance(subval, torch.Tensor) else subval for subval in batch[k]] else: batch[k] = batch[k] # motion = motion.to(self.device) # cond['y'] = {key: val.to(self.device) if torch.is_tensor(val) else val for key, val in cond['y'].items()} # st_idxes = cond['y']['st_idx'] # st_idxes # ed_idxes = cond['y']['ed_idx'] # ed_idxes # pert_verts = cond['y']['pert_verts'] # verts = cond['y']['verts'] # if 'avg_joints' in cond['y']: # avg_joints = cond['y']['avg_joints'] # std_joints = cond['y']['std_joints'] # else: # avg_joints = None # std_joints = None st_idxes = batch['st_idx'] ed_idxes = batch['ed_idx'] pert_verts = batch['pert_verts'] verts = batch['verts'] # tot pert verts tot_pert_verts.append(pert_verts) tot_verts.append(verts) ## generative denoising -> we want to use it for the denoising task ## # std_joints: bsz x 1 # avg_joints: bsz x 1 x 3 --> mean of joints for each batch ## predict_single_step ## # model_output, target = self.predict_single_step(batch, use_t=1) ### trainingjloop ours use_t = self.args.use_t tot_pred_outputs = self.predict_single_step(batch, use_t=use_t) #### diff baes jts e ## if len(tot_pred_outputs) == 3: model_output, target, pred_joints_quant = tot_pred_outputs tot_pred_joints_quant.append(pred_joints_quant) elif self.args.diff_basejtse: model_output, target, dec_disp_e_along_normals, dec_disp_e_vt_normals = tot_pred_outputs else: model_output, target = tot_pred_outputs[:2] # model output; target # ## model_output: ([6, 21, 3, 60]), target: torch.Size([6, 21, 3, 60]) # if avg_joints is not None: # ### model_output, target ### # model_output = (model_output * std_joints.unsqueeze(-1).unsqueeze(-1)) + avg_joints.unsqueeze(-1) # target = (target * std_joints.unsqueeze(-1).unsqueeze(-1)) + avg_joints.unsqueeze(-1) # 10 -> the output sequence is still a little bit noisy # # 100 -> 60 # the difficulty of predicting base pts rel position information # # the difficulty of the prediction problem # base pts rel information p ## predicting base pts relative positions to the base_pts predictions ## ### base pts predictions ## wu le ## if self.args.diff_basejtse: tot_dec_disp_e_along_normals.append(dec_disp_e_along_normals) tot_dec_disp_e_vt_normals.append(dec_disp_e_vt_normals) tot_st_idxes.append(st_idxes) tot_ed_idxes.append(ed_idxes) tot_targets.append(target) tot_model_outputs.append(model_output) # tot_model_outputs.extend(model_output) # tot_model_outputs = tot_model_outputs + model_output tot_st_idxes = torch.cat(tot_st_idxes, dim=0) tot_ed_idxes = torch.cat(tot_ed_idxes, dim=0) tot_targets = torch.cat(tot_targets, dim=0) tot_model_outputs = torch.cat(tot_model_outputs, dim=0) if self.args.diff_basejtse: tot_dec_disp_e_along_normals = torch.cat(tot_dec_disp_e_along_normals, dim=0) tot_dec_disp_e_vt_normals = torch.cat(tot_dec_disp_e_vt_normals, dim=0) if len(tot_pred_joints_quant) > 0: tot_pred_joints_quant = torch.cat(tot_pred_joints_quant, dim=0) tot_pert_verts = torch.cat(tot_pert_verts, dim=0) tot_verts = torch.cat(tot_verts, dim=0) if isinstance(tot_pred_joints_quant, torch.Tensor): return tot_targets, tot_model_outputs, tot_st_idxes, tot_ed_idxes, tot_pert_verts, tot_verts, tot_pred_joints_quant elif self.args.diff_basejtse: return tot_targets, tot_model_outputs, tot_st_idxes, tot_ed_idxes, tot_pert_verts, tot_verts, tot_dec_disp_e_along_normals, tot_dec_disp_e_vt_normals else: return tot_targets, tot_model_outputs, tot_st_idxes, tot_ed_idxes, tot_pert_verts, tot_verts def _anneal_lr(self): if not self.lr_anneal_steps: return frac_done = (self.step + self.resume_step) / self.lr_anneal_steps lr = self.lr * (1 - frac_done) for param_group in self.opt.param_groups: param_group["lr"] = lr def log_step(self): logger.logkv("step", self.step + self.resume_step) logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) def ckpt_file_name(self): return f"model{(self.step+self.resume_step):09d}.pt" def save(self): def save_checkpoint(params): if self.args.finetune_with_cond: # state_dict = self.mp_trainer.model.state_dict() else: state_dict = self.mp_trainer.master_params_to_state_dict(params) # Do not save CLIP weights clip_weights = [e for e in state_dict.keys() if e.startswith('clip_model.')] for e in clip_weights: del state_dict[e] # logger.log(f"saving model...") filename = self.ckpt_file_name() model_sv_fn = os.path.join(self.save_dir, filename) logger.log(f"saving model to {model_sv_fn}...") with bf.BlobFile(bf.join(self.save_dir, filename), "wb") as f: torch.save(state_dict, f) save_checkpoint(self.mp_trainer.master_params) with bf.BlobFile( bf.join(self.save_dir, f"opt{(self.step+self.resume_step):09d}.pt"), "wb", ) as f: torch.save(self.opt.state_dict(), f) def parse_resume_step_from_filename(filename): """ Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the checkpoint's number of steps. """ split = filename.split("model") if len(split) < 2: return 0 split1 = split[-1].split(".")[0] try: return int(split1) except ValueError: return 0 def get_blob_logdir(): # You can change this to be a separate path to save checkpoints to # a blobstore or some external drive. return logger.get_dir() def find_resume_checkpoint(): # On your infrastructure, you may want to override this to automatically # discover the latest checkpoint on your blob storage, etc. return None def log_loss_dict(diffusion, ts, losses): for key, values in losses.items(): logger.logkv_mean(key, values.mean().item()) # Log the quantiles (four quartiles, in particular). for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): quartile = int(4 * sub_t / diffusion.num_timesteps) logger.logkv_mean(f"{key}_q{quartile}", sub_loss)