From 5d9775735d9fde87e8ea397337b8fdb0b479e0f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B4=B9=E6=94=BF=E8=81=AA?= <3300949577@qq.com> Date: Wed, 14 Aug 2024 14:22:50 +0800 Subject: [PATCH] debug --- .../rectified_flow.cpython-310.pyc | Bin 2869 -> 2997 bytes diffusion/rectified_flow.py | 17 ++++++---- train.py | 32 +++++++++++++----- 3 files changed, 34 insertions(+), 15 deletions(-) diff --git a/diffusion/__pycache__/rectified_flow.cpython-310.pyc b/diffusion/__pycache__/rectified_flow.cpython-310.pyc index a44871b88477465058870e7f4e918978aa322029..4d189c389f64e73bc9596bf3564279bbba604b03 100644 GIT binary patch delta 1121 zcmb7^y>HV%6u^D{I6EeB`Vrc+1;*03J^)UWvu1 zRZ${8q%_fX(aY_zKI`+Oov_}RzsM%!GFHJ=W3S~#D%nGt&zYz!pbWdCJjowMCW*Vc z5t1F7xLNi^8Am$%t}G*s74S2(#@^ydRZlkch-FhlL#?*(A=O50y)Au|US--st!oX~ zZ@jwO2kZ88tdkswCpkW4$S}ww95If(fQE#e2F)437+@T59xwqo+%qZgX^ibg!wX!~ zAQ`Y;2V4U~Drow?<%EG7r>R|Uumjb2It3QdjbdD_+nXMZQI*Ujx|T!cxGDIEqamnDtnbvNgd{Ugp|5cnZwv{UE-r}*zZ4NVgEE#ea zO!okb07!!@0qz4H03Z#r41heyL%<~9*nwud*;qT(Xv=!ppb(lujfei7S3#e|HzC#n zbM(S4^eUMLB-o6u@VuyR&WC?E(TWk;K&qu_OYz#-JcUMG%d|#E%1hlxACSmu|bxEFXt8 z9MF>|oSPZ~M-$`4koX69(SN|C|AALe9(}JHiP^-1ZTicb_vXEs_tAr+2YptQ%b5h8 zm%lzPi|mdyaHa?=k9ag7i)4;?%+p>FPxDEYR@t_;sCo1*u~R?LX`#H+)wA<@y~JRl z!6%D=2H7Tibe~Ytq)p3K zI>4kHUX!XBc8z&(*iu@mq5@92f`VFEq1cgh+aq2_t7kxSp6t)lER0bDVPEN?G=y56)j#FmMYFm^l2Ia4d9UQJwSSS<3r* oXkxQdt>`QaV}3jG)&f7#i;iRz+OpCoYsr==3ehB9P$vcU8(a0wZvX%Q diff --git a/diffusion/rectified_flow.py b/diffusion/rectified_flow.py index 1ef960e..ac22743 100644 --- a/diffusion/rectified_flow.py +++ b/diffusion/rectified_flow.py @@ -1,11 +1,16 @@ import torch +from torch.nn.parallel import DistributedDataParallel as DDP class RectifiedFlow(torch.nn.Module): def __init__(self, model, ln=True): super().__init__() self.model = model self.ln = ln - self.stratified = False + self.stratified = False + if isinstance(model, DDP): + self.learn_sigma = model.module.learn_sigma + else: + self.learn_sigma = model.learn_sigma def forward(self, x, cond): @@ -31,7 +36,7 @@ class RectifiedFlow(torch.nn.Module): # make t, zt into same dtype as x zt, t = zt.to(x.dtype), t.to(x.dtype) vtheta = self.model(zt, t, cond) - if self.model.learn_sigma == True: + if self.learn_sigma == True: vtheta, _ = vtheta.chunk(2, dim=1) batchwise_mse = ((z1 - x - vtheta) ** 2).mean(dim=list(range(1, len(x.shape)))) tlist = batchwise_mse.detach().cpu().reshape(-1).tolist() @@ -49,11 +54,11 @@ class RectifiedFlow(torch.nn.Module): t = torch.tensor([t] * b).to(z.device) vc = self.model(z, t, cond) - if self.model.learn_sigma == True: + if self.learn_sigma == True: vc, _ = vc.chunk(2, dim=1) if null_cond is not None: vu = self.model(z, t, null_cond) - if self.model.learn_sigma == True: + if self.learn_sigma == True: vu, _ = vu.chunk(2, dim=1) vc = vu + cfg * (vc - vu) @@ -72,11 +77,11 @@ class RectifiedFlow(torch.nn.Module): t = torch.tensor([t] * b).to(z.device) vc = self.model(z, t, cond) - if self.model.learn_sigma == True: + if self.learn_sigma == True: vc, _ = vc.chunk(2, dim=1) if null_cond is not None: vu = self.model(z, t, null_cond) - if self.model.learn_sigma == True: + if self.learn_sigma == True: vu, _ = vu.chunk(2, dim=1) vc = vu + cfg * (vc - vu) x = z - i * dt * vc diff --git a/train.py b/train.py index 44b3bf3..1214bf3 100644 --- a/train.py +++ b/train.py @@ -28,7 +28,8 @@ import logging import os from models import DiT_models -from diffusion import create_diffusion +from diffusion import create_diffusion +from diffusion.rectified_flow import RectifiedFlow from diffusers.models import AutoencoderKL from download import find_model @@ -128,7 +129,7 @@ def main(args): os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders) experiment_index = len(glob(f"{args.results_dir}/*")) model_string_name = args.model.replace("/", "-") # e.g., DiT-XL/2 --> DiT-XL-2 (for naming folders) - experiment_dir = f"{args.results_dir}/{experiment_index:03d}-{model_string_name}" # Create an experiment folder + experiment_dir = f"{args.results_dir}/{model_string_name}-{args.num_experts}E{args.num_experts_per_tok}A{args.pretraining_tp}TP" # Create an experiment folder checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints os.makedirs(checkpoint_dir, exist_ok=True) logger = create_logger(experiment_dir) @@ -157,7 +158,13 @@ def main(args): requires_grad(ema, False) # model = DDP(model.to(device), device_ids=[rank]) model = DDP(model.to(device), device_ids=[device]) - diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule + + if args.rf: + logger.info("train with rectified flow") + diffusion = RectifiedFlow(model) + else: + diffusion = create_diffusion(timestep_respacing="") # default: 1000 steps, linear noise schedule + vae = AutoencoderKL.from_pretrained(args.vae_path).to(device) logger.info(f"DiT Parameters: {sum(p.numel() for p in model.parameters()):,}") @@ -212,10 +219,15 @@ def main(args): with torch.no_grad(): # Map input images to latent space + normalize latents: x = vae.encode(x).latent_dist.sample().mul_(0.18215) - t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) - model_kwargs = dict(y=y) - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() + + if args.rf: + loss, _ = diffusion.forward(x, y) + else: + t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) + model_kwargs = dict(y=y) + loss_dict = diffusion.training_losses(model, x, t, model_kwargs) + loss = loss_dict["loss"].mean() + if (data_iter_step + 1) % args.accum_iter == 0: opt.zero_grad() loss.backward() @@ -278,9 +290,11 @@ if __name__ == "__main__": parser.add_argument("--global-seed", type=int, default=2024) parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--log-every", type=int, default=100) - parser.add_argument('--accum_iter', default=8, type=int,) + parser.add_argument('--accum_iter', default=1, type=int,) parser.add_argument('--num_experts', default=8, type=int,) parser.add_argument('--num_experts_per_tok', default=2, type=int,) - parser.add_argument("--ckpt-every", type=int, default=50_000) + parser.add_argument('--pretraining_tp', default=2, type=int,) + parser.add_argument("--ckpt-every", type=int, default=50_000) + parser.add_argument("--rf", type=bool, default=False) args = parser.parse_args() main(args)