From c5cc1318c2a7597fe62d3379e50187d0b0f22538 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Thu, 1 Dec 2022 13:45:21 +0100 Subject: Update --- dreambooth.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) (limited to 'dreambooth.py') diff --git a/dreambooth.py b/dreambooth.py index 3dd0920..31dbea2 100644 --- a/dreambooth.py +++ b/dreambooth.py @@ -32,6 +32,7 @@ logger = get_logger(__name__) torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.benchmark = True def parse_args(): @@ -474,7 +475,6 @@ class Checkpointer: scheduler=self.scheduler, ).to(self.accelerator.device) pipeline.set_progress_bar_config(dynamic_ncols=True) - pipeline.enable_vae_slicing() train_data = self.datamodule.train_dataloader() val_data = self.datamodule.val_dataloader() @@ -550,6 +550,12 @@ class Checkpointer: def main(): args = parse_args() + if args.train_text_encoder and args.gradient_accumulation_steps > 1 and accelerator.num_processes > 1: + raise ValueError( + "Gradient accumulation is not supported when training the text encoder in distributed training. " + "Please set gradient_accumulation_steps to 1. This feature will be supported in the future." + ) + instance_identifier = args.instance_identifier if len(args.placeholder_token) != 0: @@ -587,6 +593,7 @@ def main(): checkpoint_scheduler = DPMSolverMultistepScheduler.from_pretrained( args.pretrained_model_name_or_path, subfolder='scheduler') + vae.enable_slicing() unet.set_use_memory_efficient_attention_xformers(True) if args.gradient_checkpointing: @@ -903,7 +910,7 @@ def main(): sample_checkpoint = False for step, batch in enumerate(train_dataloader): - with accelerator.accumulate(itertools.chain(unet, text_encoder)): + with accelerator.accumulate(unet): # Convert images to latent space latents = vae.encode(batch["pixel_values"]).latent_dist.sample() latents = latents * 0.18215 -- cgit v1.2.3-54-g00ecf