From 784abf3d6a18bd6f8556746026ef4d1dccbd8096 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 09:43:18 +0800 Subject: [PATCH] support rectified flow --- .../rectified_flow.cpython-310.pyc | Bin 0 -> 2869 bytes diffusion/rectified_flow.py | 85 +++++++++ sample.py | 45 +++-- sample_ddp.py | 166 ++++++++++++++++++ test.py | 4 +- train.py | 2 +- train_deepspeed.py | 46 +++-- 7 files changed, 317 insertions(+), 31 deletions(-) create mode 100644 diffusion/__pycache__/rectified_flow.cpython-310.pyc create mode 100644 diffusion/rectified_flow.py create mode 100644 sample_ddp.py diff --git a/diffusion/__pycache__/rectified_flow.cpython-310.pyc b/diffusion/__pycache__/rectified_flow.cpython-310.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a44871b88477465058870e7f4e918978aa322029 GIT binary patch literal 2869 zcmd5;&5s+m6(>0#8jU`71ceh4v0YZUK9DY1LzW4h*a?)yr z4CSkT|9zmEjQxW?&KDOScTvl?AtaMLXMOe&lR`Srm~>Kh%#X#=={u4?U_JK*X0fA$ z+>kHZXnhkiRsgCuo^q#T4n|dpM&175FW%!IQdB1kn(~!-oC-|DV ztY61|JGS*{n&@nhrSg70IenNo)jk_RZGlSvmHnD6*@~|jXNDUV2_qn3IixEY^ogVU zJ+?L{U~`I*P0}p2#$}A}?%MY87ub>^asg(|n*j=vO|(b8hQVl;m+)M$vQj zpU)4*aaPD!#|M=tN3lxf!FZPI3@YRNU_MHVVKOO(`-3z)NU+U;%=C}PlXUwr7|hCS zQXHu5Iimp%xUpRw1RtW8yFC2q_ zsv}6D4z;s#B?4p|DXy`#v+@i-_A$#faF=wx64x2#)QmI0+U=-uZnA4^;bA0@uJJkm zOP8&=uCGGru6buO)-WD4`Na*)YhKi_sG+Z6YSno7|2r&vZ>!>%R17F=-vqf9hWblmXTiC0;6!&2JVh1y?0{(TglktH7B*PTAv5ez08{xw0P$iV{w$tAscK-b62P$7f6oei()V;e zDa%(wln$5;@8nsjlgW5`_ztP1&AV*gyNI&EM4iS;?%&BLNt~C5Xlz&Y+h4>w8J%Wj z8qo&5x4@cgo>qiWHQK||6P1i?kY`0Xjgu7SOsp@{lMHQF#f7wfss>qcV%>49M^NxY zX(P|Lx+htBYMneSz~E6lO>I~WM0xf&wVu>pg4M)T+sM;c6%nnEt(S~u z#banrW9T5eD}}?9Owch5bCT)I)>OJt>g!20$(d)<-mW4ZS^m_zV6n9P3tN9Oiwm9Q zX=(Y~a&1MSt*g_gQ!D20qw2nYqDN3|dDL&c^3J$SDF`s2@9WnFqxRCzl(NeoL4<-> zf5W-WZ;2-FRP{EpKnHzo-sD&KyL=B{yL^82wGlgd@e`=|Cn}Y@31zWuX`^wLC7S26 zgbzl!HyL=!SHjeW;7|t)Upe3r7!E%9;3oK~ee~3%bBUWZ!-vGs0C|KMFJOAgr&`qMRs2%v9G&yiMZ!Bz{1mN20n+en>6ioAu-AG%aMWT>(w~lm_0T zfqF5^^QZz{vy8{nJdH}7PD`M5FjVi;6d|?svvE8mC}!$|s_!IGKif2Bm$~k5^pF|l zZHQ0={6pYtmj}Gd+d$V9v^KcfcO{3Uu_=v{`Z>P-|O&dtmkgDi=wdpA1Xtt{&UJIUsbnXaA#xcAyCUox?2O* G75@WC;;jMz literal 0 HcmV?d00001 diff --git a/diffusion/rectified_flow.py b/diffusion/rectified_flow.py new file mode 100644 index 0000000..1ef960e --- /dev/null +++ b/diffusion/rectified_flow.py @@ -0,0 +1,85 @@ +import torch + +class RectifiedFlow(torch.nn.Module): + def __init__(self, model, ln=True): + super().__init__() + self.model = model + self.ln = ln + self.stratified = False + + def forward(self, x, cond): + + b = x.size(0) + if self.ln: + if self.stratified: + # stratified sampling of normals + # first stratified sample from uniform + quantiles = torch.linspace(0, 1, b + 1).to(x.device) + z = quantiles[:-1] + torch.rand((b,)).to(x.device) / b + # now transform to normal + z = torch.erfinv(2 * z - 1) * math.sqrt(2) + t = torch.sigmoid(z) + else: + nt = torch.randn((b,)).to(x.device) + t = torch.sigmoid(nt) + else: + t = torch.rand((b,)).to(x.device) + texp = t.view([b, *([1] * len(x.shape[1:]))]) + z1 = torch.randn_like(x) + zt = (1 - texp) * x + texp * z1 + + # 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: + 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() + ttloss = [(tv, tloss) for tv, tloss in zip(t, tlist)] + return batchwise_mse.mean(), {"batchwise_loss": ttloss} + + @torch.no_grad() + def sample(self, z, cond, null_cond=None, sample_steps=50, cfg=2.0): + b = z.size(0) + dt = 1.0 / sample_steps + dt = torch.tensor([dt] * b).to(z.device).view([b, *([1] * len(z.shape[1:]))]) + images = [z] + for i in range(sample_steps, 0, -1): + t = i / sample_steps + t = torch.tensor([t] * b).to(z.device) + + vc = self.model(z, t, cond) + if self.model.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: + vu, _ = vu.chunk(2, dim=1) + vc = vu + cfg * (vc - vu) + + z = z - dt * vc + images.append(z) + return images + + @torch.no_grad() + def sample_with_xps(self, z, cond, null_cond=None, sample_steps=50, cfg=2.0): + b = z.size(0) + dt = 1.0 / sample_steps + dt = torch.tensor([dt] * b).to(z.device).view([b, *([1] * len(z.shape[1:]))]) + images = [z] + for i in range(sample_steps, 0, -1): + t = i / sample_steps + t = torch.tensor([t] * b).to(z.device) + + vc = self.model(z, t, cond) + if self.model.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: + vu, _ = vu.chunk(2, dim=1) + vc = vu + cfg * (vc - vu) + x = z - i * dt * vc + z = z - dt * vc + images.append(x) + return images \ No newline at end of file diff --git a/sample.py b/sample.py index 5b1282a..0a22458 100644 --- a/sample.py +++ b/sample.py @@ -15,6 +15,7 @@ from diffusion import create_diffusion from diffusers.models import AutoencoderKL from download import find_model from models import DiT_models +from diffusion.rectified_flow import RectifiedFlow import argparse @@ -60,17 +61,21 @@ def main(args): elif args.model == "DiT-B/2": ckpt_path = "dit_moe_b_8E2A.pt" elif args.model == "DiT-XL/2": - ckpt_path = "results/deepspeed-DiT-XL-2/checkpoints/ckpt.pt" + ckpt_path = "results/deepspeed-DiT-XL-2-rf/checkpoints/tmp.pt" else: - ckpt_path = "results/deepspeed-DiT-G-2/checkpoints/ckpt.pt" + ckpt_path = "results/deepspeed-DiT-G-2-rf/checkpoints/tmp.pt" else: ckpt_path = args.ckpt state_dict = find_model(ckpt_path) model.load_state_dict(state_dict) - model.eval() # important! - diffusion = create_diffusion(str(args.num_sampling_steps)) + model.eval() # important! + + if args.rf: + diffusion = RectifiedFlow(model) + else: + diffusion = create_diffusion(str(args.num_sampling_steps)) vae = AutoencoderKL.from_pretrained(args.vae_path).to(device) # Labels to condition the model with (feel free to change): @@ -88,17 +93,28 @@ def main(args): model_kwargs = dict(y=y, cfg_scale=args.cfg_scale) if dtype == torch.float16: - with torch.autocast(device_type='cuda'): - samples = diffusion.p_sample_loop( - model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device - ) + if args.rf: + with torch.autocast(device_type='cuda'): + STEPSIZE = 50 + init_noise = torch.randn(n, 4, latent_size, latent_size, device=device) + conds = torch.tensor(class_labels, device=device) + images = diffusion.sample_with_xps(init_noise, conds, null_cond = torch.tensor([1000] * n).cuda(), sample_steps = STEPSIZE, cfg = 7.0) + samples = vae.decode(images[-1] / 0.18215).sample + + else: + with torch.autocast(device_type='cuda'): + samples = diffusion.p_sample_loop( + model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device + ) + samples, _ = samples.chunk(2, dim=0) # Remove null class samples + samples = vae.decode(samples / 0.18215).sample + else: samples = diffusion.p_sample_loop( model.forward_with_cfg, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device ) - - samples, _ = samples.chunk(2, dim=0) # Remove null class samples - samples = vae.decode(samples / 0.18215).sample + samples, _ = samples.chunk(2, dim=0) # Remove null class samples + samples = vae.decode(samples / 0.18215).sample # Save and display images: if args.model == "DiT-S/2": @@ -113,15 +129,16 @@ def main(args): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2") + parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-G/2") parser.add_argument("--vae-path", type=str, default="/maindata/data/shared/multimodal/zhengcong.fei/ckpts/sd-vae-ft-mse") parser.add_argument("--image-size", type=int, choices=[256, 512], default=256) parser.add_argument("--num-classes", type=int, default=1000) parser.add_argument("--cfg-scale", type=float, default=4.0) - parser.add_argument('--num_experts', default=8, type=int,) + parser.add_argument('--num_experts', default=16, type=int,) parser.add_argument('--num_experts_per_tok', default=2, type=int,) parser.add_argument("--num-sampling-steps", type=int, default=250) parser.add_argument("--seed", type=int, default=2024) parser.add_argument("--ckpt", type=str, default=None, ) + parser.add_argument("--rf", type=bool, default=True) args = parser.parse_args() - main(args) + main(args) diff --git a/sample_ddp.py b/sample_ddp.py new file mode 100644 index 0000000..0a6b1ab --- /dev/null +++ b/sample_ddp.py @@ -0,0 +1,166 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +""" +Samples a large number of images from a pre-trained DiT model using DDP. +Subsequently saves a .npz file that can be used to compute FID and other +evaluation metrics via the ADM repo: https://github.com/openai/guided-diffusion/tree/main/evaluations + +For a simple single-GPU/CPU sampling script, see sample.py. +""" +import torch +import torch.distributed as dist +from models import DiT_models +from download import find_model +from diffusion import create_diffusion +from diffusers.models import AutoencoderKL +from tqdm import tqdm +import os +from PIL import Image +import numpy as np +import math +import argparse + + +def create_npz_from_sample_folder(sample_dir, num=50_000): + """ + Builds a single .npz file from a folder of .png samples. + """ + samples = [] + for i in tqdm(range(num), desc="Building .npz file from samples"): + sample_pil = Image.open(f"{sample_dir}/{i:06d}.png") + sample_np = np.asarray(sample_pil).astype(np.uint8) + samples.append(sample_np) + samples = np.stack(samples) + assert samples.shape == (num, samples.shape[1], samples.shape[2], 3) + npz_path = f"{sample_dir}.npz" + np.savez(npz_path, arr_0=samples) + print(f"Saved .npz file to {npz_path} [shape={samples.shape}].") + return npz_path + + +def main(args): + """ + Run sampling. + """ + torch.backends.cuda.matmul.allow_tf32 = args.tf32 # True: fast but may lead to some small numerical differences + assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage" + torch.set_grad_enabled(False) + + # Setup DDP: + dist.init_process_group("nccl") + rank = dist.get_rank() + device = rank % torch.cuda.device_count() + seed = args.global_seed * dist.get_world_size() + rank + torch.manual_seed(seed) + torch.cuda.set_device(device) + print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.") + + if args.ckpt is None: + assert args.model == "DiT-XL/2", "Only DiT-XL/2 models are available for auto-download." + assert args.image_size in [256, 512] + assert args.num_classes == 1000 + + # Load model: + latent_size = args.image_size // 8 + model = DiT_models[args.model]( + input_size=latent_size, + num_classes=args.num_classes + ).to(device) + # Auto-download a pre-trained model or load a custom DiT checkpoint from train.py: + ckpt_path = args.ckpt or f"DiT-XL-2-{args.image_size}x{args.image_size}.pt" + state_dict = find_model(ckpt_path) + model.load_state_dict(state_dict) + model.eval() # important! + diffusion = create_diffusion(str(args.num_sampling_steps)) + vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device) + assert args.cfg_scale >= 1.0, "In almost all cases, cfg_scale be >= 1.0" + using_cfg = args.cfg_scale > 1.0 + + # Create folder to save samples: + model_string_name = args.model.replace("/", "-") + ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained" + folder_name = f"{model_string_name}-{ckpt_string_name}-size-{args.image_size}-vae-{args.vae}-" \ + f"cfg-{args.cfg_scale}-seed-{args.global_seed}" + sample_folder_dir = f"{args.sample_dir}/{folder_name}" + if rank == 0: + os.makedirs(sample_folder_dir, exist_ok=True) + print(f"Saving .png samples at {sample_folder_dir}") + dist.barrier() + + # Figure out how many samples we need to generate on each GPU and how many iterations we need to run: + n = args.per_proc_batch_size + global_batch_size = n * dist.get_world_size() + # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples: + total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size) + if rank == 0: + print(f"Total number of images that will be sampled: {total_samples}") + assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size" + samples_needed_this_gpu = int(total_samples // dist.get_world_size()) + assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size" + iterations = int(samples_needed_this_gpu // n) + pbar = range(iterations) + pbar = tqdm(pbar) if rank == 0 else pbar + total = 0 + for _ in pbar: + # Sample inputs: + z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device) + y = torch.randint(0, args.num_classes, (n,), device=device) + + # Setup classifier-free guidance: + if using_cfg: + z = torch.cat([z, z], 0) + y_null = torch.tensor([1000] * n, device=device) + y = torch.cat([y, y_null], 0) + model_kwargs = dict(y=y, cfg_scale=args.cfg_scale) + sample_fn = model.forward_with_cfg + else: + model_kwargs = dict(y=y) + sample_fn = model.forward + + # Sample images: + samples = diffusion.p_sample_loop( + sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=False, device=device + ) + if using_cfg: + samples, _ = samples.chunk(2, dim=0) # Remove null class samples + + samples = vae.decode(samples / 0.18215).sample + samples = torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy() + + # Save samples to disk as individual .png files + for i, sample in enumerate(samples): + index = i * dist.get_world_size() + rank + total + Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png") + total += global_batch_size + + # Make sure all processes have finished saving their samples before attempting to convert to .npz + dist.barrier() + if rank == 0: + create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples) + print("Done.") + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, choices=list(DiT_models.keys()), default="DiT-XL/2") + parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="ema") + parser.add_argument("--sample-dir", type=str, default="samples") + parser.add_argument("--per-proc-batch-size", type=int, default=32) + parser.add_argument("--num-fid-samples", type=int, default=50_000) + parser.add_argument("--image-size", type=int, choices=[256, 512], default=256) + parser.add_argument("--num-classes", type=int, default=1000) + parser.add_argument("--cfg-scale", type=float, default=1.5) + parser.add_argument("--num-sampling-steps", type=int, default=250) + parser.add_argument("--global-seed", type=int, default=0) + parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True, + help="By default, use TF32 matmuls. This massively accelerates sampling on Ampere GPUs.") + parser.add_argument("--ckpt", type=str, default=None, + help="Optional path to a DiT checkpoint (default: auto-download a pre-trained DiT-XL/2 model).") + args = parser.parse_args() + main(args) diff --git a/test.py b/test.py index 6043d4c..473a40f 100644 --- a/test.py +++ b/test.py @@ -106,12 +106,12 @@ def image_class_expert_ratio(): def ckpts_clean(): # only save ema ckpts for ckpt uploading - ckpt_path = 'results/003-DiT-B-2/checkpoints/ckpt.pt' + ckpt_path = 'results/003-DiT-B-2/checkpoints/ckpt2.pt' checkpoint = torch.load(ckpt_path, map_location=lambda storage, loc: storage) new_checkpoint = { "ema": checkpoint['ema'], } - torch.save(new_checkpoint, 'ckpt_clean.pt') + torch.save(new_checkpoint, 'dit_moe_b_8E2A.pt') diff --git a/train.py b/train.py index 9a16c73..44b3bf3 100644 --- a/train.py +++ b/train.py @@ -275,7 +275,7 @@ if __name__ == "__main__": parser.add_argument("--num-classes", type=int, default=1000) parser.add_argument("--epochs", type=int, default=1400) parser.add_argument("--global-batch-size", type=int, default=64) - parser.add_argument("--global-seed", type=int, default=1234) + 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,) diff --git a/train_deepspeed.py b/train_deepspeed.py index f00cce8..5a32456 100644 --- a/train_deepspeed.py +++ b/train_deepspeed.py @@ -29,6 +29,7 @@ import os from models import DiT_models from diffusion import create_diffusion +from diffusion.rectified_flow import RectifiedFlow from diffusers.models import AutoencoderKL from download import find_model import deepspeed @@ -115,7 +116,10 @@ def main(args): # Setup an experiment folder model_string_name = args.model.replace("/", "-") # e.g., DiT-XL/2 --> DiT-XL-2 (for naming folders) - experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}" # Create an experiment folder + if args.rf: + experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}-rf" + else: + experiment_dir = f"{args.results_dir}/deepspeed-{model_string_name}" # Create an experiment folder checkpoint_dir = f"{experiment_dir}/checkpoints" # Stores saved model checkpoints if rank == 0: @@ -142,9 +146,12 @@ def main(args): print('load from: ', args.resume) state_dict = find_model(args.resume) model.load_state_dict(state_dict) - - - 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()):,}") @@ -194,11 +201,17 @@ 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) - with torch.autocast(device_type='cuda'): - loss_dict = diffusion.training_losses(model, x, t, model_kwargs) - loss = loss_dict["loss"].mean() + if args.rf: + with torch.autocast(device_type='cuda'): + loss, _ = diffusion.forward(x, y) + + else: + t = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=device) + model_kwargs = dict(y=y) + with torch.autocast(device_type='cuda'): + loss_dict = diffusion.training_losses(model, x, t, model_kwargs) + loss = loss_dict["loss"].mean() + model_engine.backward(loss) model_engine.step() @@ -223,9 +236,13 @@ def main(args): start_time = time() # Save DiT checkpoint: - if train_steps % args.ckpt_every == 0 and train_steps > 0: - checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}" - model_engine.save_checkpoint(checkpoint_path) + if train_steps % args.ckpt_every == 0 and train_steps > 0: + try: + checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}" + model_engine.save_checkpoint(checkpoint_path) + except Exception as e: + print(e) + dist.barrier() # model.eval() # important! This disables randomized embedding dropout @@ -244,14 +261,15 @@ if __name__ == "__main__": parser.add_argument("--num-classes", type=int, default=1000) parser.add_argument("--epochs", type=int, default=1400) parser.add_argument("--train_batch_size", type=int, default=2) - parser.add_argument("--global-seed", type=int, default=1234) - parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--global-seed", type=int, default=2024) + parser.add_argument("--num-workers", type=int, default=0) parser.add_argument("--log-every", type=int, default=100) parser.add_argument('--accum_iter', default=8, 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=10_000) parser.add_argument('--local-rank', type=int, default=-1, help='local rank passed from distributed launcher') + parser.add_argument("--rf", type=bool, default=False) parser = deepspeed.add_config_arguments(parser) args = parser.parse_args() print(args)