Updated Weights & Biases logging + Made codebase pip installable #22

soumik12345 · 2023-02-10T16:35:43Z

Changes included in this PR:

Added use of the WandbMetricsLogger callback logging training metrics to Weights & Biases.
Added use of WandbModelCheckpoint for logging model checkpoints as Weights & Biases artifacts.
Added an additional callback QualitativeValidationCallback for visualizing the generations corresponding to multiple prompts in an epoch-wise manner using Weights & Biases Table.
Updated the validation_prompt argument in the training script to accept multiple prompts for visualizing the generations during training.
Made the codebase pip installable

Todos:

Make sure to log the results corresponding to all prompts in both table and media panel.
Update README to reflect all the changes.

Sample Weights & Biases run: https://wandb.ai/geekyrakshit/dreambooth-keras/runs/h0kdq2dw

sayakpaul · 2023-02-10T16:49:07Z

@soumik12345 thanks for your PR!

I think we would want to keep the codebase a bit more lightweight for people to be able to navigate around easily and customize as per their needs.

So, with that in mind, I would request to introduce the changes only pertaining to WandB logging (images). Also, please keep the code-lints unchanged.

Here's what we can do I think (@deep-diver feel free to chime in):

Keep your logging (image) related changes. We will mention you from the repository crediting your work (of course).
You package this repository from a separate repository (with due credits).

train_dreambooth.py

…ompts

soumik12345 · 2023-02-10T20:11:34Z

Hi @sayakpaul
I updated the PR as per your request. Here's a sample wandb run: https://wandb.ai/geekyrakshit/dreambooth-keras/runs/4wydmf8w

I will edit the readme to include the changes, once you approve the changes.

sayakpaul · 2023-02-11T14:30:51Z

Hi @sayakpaul I updated the PR as per your request. Here's a sample wandb run: https://wandb.ai/geekyrakshit/dreambooth-keras/runs/4wydmf8w

I will edit the readme to include the changes, once you approve the changes.

I still don't see the prompt in the validation panel (not the table but the panel).

@soumik12345

soumik12345 · 2023-02-13T10:44:55Z

Hi @sayakpaul, I updated the logging of generated images from validation prompt on W&B panel. Here's a sample run: https://wandb.ai/geekyrakshit/dreambooth-keras/runs/8shvurfp

sayakpaul · 2023-02-14T07:08:27Z

src/utils.py

+    for prompt in tqdm(prompts):
+        images_dreamboothed = sd_model.text_to_image(prompt, batch_size=num_imgs_to_gen)
+        wandb.log(
+            {
+                f"validation/Prompt: {prompt}": [
+                    wandb.Image(PIL.Image.fromarray(image), caption=f"{i}: {prompt}")
+                    for i, image in enumerate(images_dreamboothed)
+                ]
+            }
+        )


sayakpaul · 2023-02-14T07:12:49Z

train_dreambooth.py

+                num_imgs_to_gen=args.num_images_to_generate,
+            ),
+        ]
+        if args.log_wandb


Please refactor it such that the else block becomes more pronounced. No comprehensions here, please.

Not addressed.

@sayakpaul made the change.

sayakpaul · 2023-02-14T07:14:36Z

train_dreambooth.py

+        else [
+            tf.keras.callbacks.ModelCheckpoint(
+                ckpt_path_prefix, save_weights_only=True, monitor="loss", mode="min"
            )
+        ]
+    )


The problem with only saving the low-loss checkpoints in DreamBooth is that the loss curves are not that informative if you take a look at them (think of vanilla GAN training loss curves).

So, this is why I just logged weights after the end of training.

In my experience, the low-loss checkpoints often lead to inferior results.

I'll remove the monitor in that case.

sayakpaul · 2023-02-14T07:14:51Z

train_dreambooth.py

+            DreamBoothCheckpointCallback(
+                ckpt_path_prefix, save_weights_only=True, monitor="loss", mode="min"
+            ),


Same as https://github.com/sayakpaul/dreambooth-keras/pull/22/files#r1105388125

sayakpaul

Thanks for the changes.

sayakpaul

One change remains.

sayakpaul · 2023-02-14T11:54:57Z

train_dreambooth.py

+    if args.log_wandb:
+        # log training metrics to Weights & Biases
+        callbacks.append(WandbMetricsLogger(log_freq="batch"))
+        # log model checkpoints to Weights & Biases as artifacts
+        callbacks.append(
+            DreamBoothCheckpointCallback(ckpt_path_prefix, save_weights_only=True)
+        )
+        # perform inference on validation prompts at the end of every epoch and
+        # log the resuts to a Weights & Biases table
+        callbacks.append(
            QualitativeValidationCallback(
                img_heigth=args.img_resolution,
                img_width=args.img_resolution,
                prompts=validation_prompts,
                num_imgs_to_gen=args.num_images_to_generate,


Perfect 💯

sayakpaul · 2023-02-14T11:55:31Z

LGTM! @deep-diver what do you think?

deep-diver · 2023-02-14T22:50:59Z

train_dreambooth.py

+    else:
+        callbacks.append(
+            tf.keras.callbacks.ModelCheckpoint(ckpt_path_prefix, save_weights_only=True)
+        )


DreamBoothCheckpointCallback is a subclass of
WandbModelCheckpoint which ultimately inherits tf.keras.callbacks.ModelCheckpoint.

I think we don't need this separate else branch. Instead, we could simply pass args.log_wandb to DreamBoothCheckpointCallback to decide if ckpts should be uploaded to W&B.

deep-diver · 2023-02-14T22:52:56Z

train_dreambooth.py

+    validation_prompts = [f"A photo of {args.unique_id} {args.class_category} in a bucket"]
+    if args.validation_prompts is not None:
+        validation_prompts += args.validation_prompts


nit: if validation prompts are provided, the default prompt might not make sense since prompts depend on different situations (maybe = instead of +=?)

deep-diver · 2023-02-14T23:22:09Z

train_dreambooth.py


+    train(dreambooth_trainer, train_dataset, args.max_train_steps, callbacks)
+
+    if args.log_wandb:


nit: maybe it is better to reuse StableDiffusion from QualitativeValidationCallback since it is always there under args.log_wandb=True condition.

… artifacts

…ualitativeValidationCallback + updated readme to reflect changes

deep-diver · 2023-02-15T14:43:10Z

LGTM! Thanks for your contribution!

sayakpaul · 2023-02-15T14:45:14Z

Feel free to add a note of these changes to README, @soumik12345! Thank you.

soumik12345 · 2023-02-15T15:50:27Z

Feel free to add a note of these changes to README, @soumik12345! Thank you.

@sayakpaul updated the README to reflect the changes.

soumik12345 added 12 commits February 9, 2023 11:01

update: wandb integration for model checkpoint and metrics logging

0bdaf86

fix: wandb argparse argument

ff666fd

add: QualitativeValidationCallback

b85069d

update: CLI + callback

27c3160

update: removed tqdm pbar from QualitativeValidationCallback

995c0f3

update: train_dreambooth.py

305b516

add: converted to pip installable package

140fceb

update: model checkpointing

f93f893

update: CLI + README

6dd10ff

add: random seed for reproducibility

fa4ab6a

update: readme

064c70f

refactor: codebase -> module dreambooth_keras + added packaging info

a1ad671

soumik12345 marked this pull request as draft February 10, 2023 16:35

sayakpaul reviewed Feb 10, 2023

View reviewed changes

train_dreambooth.py Outdated Show resolved Hide resolved

soumik12345 added 8 commits February 10, 2023 17:15

refactor: reverted packaging

a0a3b70

refactor: revert linting changes

6566fec

refactor: revert linting changes

9d46f44

update: made media panel log generated imaged corresponding to all pr…

ad177bd

…ompts

fix: calling log_images

d50f7c1

fix: minor logging

2c70aaf

fix: logging generated image to media panel

522b1e4

fix: logging generated image to media panel

1a047cf

soumik12345 requested a review from sayakpaul February 10, 2023 20:11

soumik12345 marked this pull request as ready for review February 10, 2023 20:11

updated image logging in W&B panels

831047d

added DreamBoothCheckpointCallback to fix model checkpoint logging

e3cb548

sayakpaul reviewed Feb 14, 2023

View reviewed changes

removed monitors from ModelCheckpoint callbacks

be49d05

soumik12345 requested a review from sayakpaul February 14, 2023 08:03

sayakpaul reviewed Feb 14, 2023

View reviewed changes

updated callback definitions in the training script

535486c

soumik12345 requested a review from sayakpaul February 14, 2023 10:59

sayakpaul reviewed Feb 14, 2023

View reviewed changes

sayakpaul approved these changes Feb 14, 2023

View reviewed changes

sayakpaul requested a review from deep-diver February 14, 2023 11:55

deep-diver requested changes Feb 14, 2023

View reviewed changes

soumik12345 added 2 commits February 15, 2023 09:48

updated DreamBoothCheckpointCallback to optionally log checkpoints as…

606dcdf

… artifacts

arefactore loogic for logging generated images to W&B panels inside Q…

f9b2543

…ualitativeValidationCallback + updated readme to reflect changes

soumik12345 requested a review from deep-diver February 15, 2023 10:52

deep-diver approved these changes Feb 15, 2023

View reviewed changes

updated readme to reflect the changes

458864f

sayakpaul approved these changes Feb 15, 2023

View reviewed changes

sayakpaul merged commit f7be949 into sayakpaul:main Feb 15, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Updated Weights & Biases logging + Made codebase pip installable #22

Updated Weights & Biases logging + Made codebase pip installable #22

soumik12345 commented Feb 10, 2023 •

edited

sayakpaul commented Feb 10, 2023 •

edited

soumik12345 commented Feb 10, 2023 •

edited

sayakpaul commented Feb 11, 2023 •

edited

soumik12345 commented Feb 13, 2023

sayakpaul Feb 14, 2023

sayakpaul Feb 14, 2023

sayakpaul Feb 14, 2023

soumik12345 Feb 14, 2023

soumik12345 Feb 14, 2023

sayakpaul Feb 14, 2023

sayakpaul Feb 14, 2023

soumik12345 Feb 14, 2023

sayakpaul Feb 14, 2023

sayakpaul left a comment

sayakpaul left a comment

sayakpaul Feb 14, 2023

sayakpaul commented Feb 14, 2023

deep-diver Feb 14, 2023

deep-diver Feb 14, 2023

deep-diver Feb 14, 2023

deep-diver commented Feb 15, 2023

sayakpaul commented Feb 15, 2023

soumik12345 commented Feb 15, 2023


		train(dreambooth_trainer, train_dataset, args.max_train_steps, callbacks)

		if args.log_wandb:

Updated Weights & Biases logging + Made codebase pip installable #22

Updated Weights & Biases logging + Made codebase pip installable #22

Conversation

soumik12345 commented Feb 10, 2023 • edited

sayakpaul commented Feb 10, 2023 • edited

soumik12345 commented Feb 10, 2023 • edited

sayakpaul commented Feb 11, 2023 • edited

soumik12345 commented Feb 13, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

sayakpaul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Feb 14, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

deep-diver commented Feb 15, 2023

sayakpaul commented Feb 15, 2023

soumik12345 commented Feb 15, 2023

soumik12345 commented Feb 10, 2023 •

edited

sayakpaul commented Feb 10, 2023 •

edited

soumik12345 commented Feb 10, 2023 •

edited

sayakpaul commented Feb 11, 2023 •

edited