From e17d5160464617841a3d266792b4f82aad6c0919 Mon Sep 17 00:00:00 2001 From: Ben Evans Date: Wed, 8 Feb 2023 12:10:22 +0000 Subject: [PATCH 1/5] Log Unconditional Image Generation Samples to WandB --- .../unconditional_image_generation/train_unconditional.py | 6 ++++++ .../unconditional_image_generation/train_unconditional.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index 586ed8c8d397..912c404c0573 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -562,6 +562,12 @@ def transform_images(examples): accelerator.get_tracker("tensorboard").add_images( "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) + elif args.logger == "wandb": + import wandb + + accelerator.get_tracker("wandb").log( + {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step + ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: # save the model diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 3a188c926d5d..34a33b0afbdf 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -609,6 +609,12 @@ def transform_images(examples): accelerator.get_tracker("tensorboard").add_images( "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) + elif args.logger == "wandb": + import wandb + + accelerator.get_tracker("wandb").log( + {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step + ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: # save the model From 8ad1a811d377592a50c50d6b7914324bedd5aba9 Mon Sep 17 00:00:00 2001 From: Ben Evans Date: Thu, 9 Feb 2023 09:58:41 +0000 Subject: [PATCH 2/5] Check for wandb installation and parity between onnxruntime script --- .../train_unconditional.py | 14 +++++++++----- .../train_unconditional.py | 10 ++++++---- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index 912c404c0573..1eef289ac01a 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -21,7 +21,7 @@ from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel from diffusers.optimization import get_scheduler from diffusers.training_utils import EMAModel -from diffusers.utils import check_min_version +from diffusers.utils import check_min_version, is_tensorboard_available, is_wandb_available # Will error if the minimal version of diffusers is not installed. Remove at your own risks. @@ -220,6 +220,7 @@ def parse_args(): help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.", ) parser.add_argument("--ddpm_num_steps", type=int, default=1000) + parser.add_argument("--ddpm_num_inference_steps", type=int, default=1000) parser.add_argument("--ddpm_beta_schedule", type=str, default="linear") parser.add_argument( "--checkpointing_steps", @@ -271,6 +272,11 @@ def main(args): logging_dir=logging_dir, ) + if args.logger == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + import wandb + # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", @@ -552,19 +558,17 @@ def transform_images(examples): generator=generator, batch_size=args.eval_batch_size, output_type="numpy", - num_inference_steps=args.ddpm_num_steps, + num_inference_steps=args.ddpm_num_inference_steps, ).images # denormalize the images and save to tensorboard images_processed = (images * 255).round().astype("uint8") - if args.logger == "tensorboard": + if args.logger == "tensorboard" and is_tensorboard_available(): accelerator.get_tracker("tensorboard").add_images( "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) elif args.logger == "wandb": - import wandb - accelerator.get_tracker("wandb").log( {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step ) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 34a33b0afbdf..9170254bd1e6 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -22,7 +22,7 @@ from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel from diffusers.optimization import get_scheduler from diffusers.training_utils import EMAModel -from diffusers.utils import check_min_version, is_tensorboard_available +from diffusers.utils import check_min_version, is_tensorboard_available, is_wandb_available # Will error if the minimal version of diffusers is not installed. Remove at your own risks. @@ -280,6 +280,11 @@ def main(args): logging_dir=logging_dir, ) + if args.logger == "wandb": + if not is_wandb_available(): + raise ImportError("Make sure to install wandb if you want to use it for logging during training.") + import wandb + # `accelerate` 0.16.0 will have better support for customized saving if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format @@ -599,7 +604,6 @@ def transform_images(examples): batch_size=args.eval_batch_size, num_inference_steps=args.ddpm_num_inference_steps, output_type="numpy", - num_inference_steps=args.ddpm_num_steps, ).images # denormalize the images and save to tensorboard @@ -610,8 +614,6 @@ def transform_images(examples): "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) elif args.logger == "wandb": - import wandb - accelerator.get_tracker("wandb").log( {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step ) From b580a90c75e94db2ea65e68a59214f78be0be59e Mon Sep 17 00:00:00 2001 From: Ben Evans Date: Tue, 14 Feb 2023 12:46:51 +0000 Subject: [PATCH 3/5] Log epoch to wandb --- .../unconditional_image_generation/train_unconditional.py | 2 +- examples/unconditional_image_generation/train_unconditional.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index 1eef289ac01a..d6d8967b43ce 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -570,7 +570,7 @@ def transform_images(examples): ) elif args.logger == "wandb": accelerator.get_tracker("wandb").log( - {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step + {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, step=global_step ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 9170254bd1e6..8bdf6fc20dfd 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -615,7 +615,7 @@ def transform_images(examples): ) elif args.logger == "wandb": accelerator.get_tracker("wandb").log( - {"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step + {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, step=global_step ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: From c8c5799952a63704b587050512e84873e2c1d2bd Mon Sep 17 00:00:00 2001 From: Ben Evans Date: Tue, 14 Feb 2023 12:50:41 +0000 Subject: [PATCH 4/5] Check for tensorboard logger early on --- .../unconditional_image_generation/train_unconditional.py | 8 ++++++-- .../unconditional_image_generation/train_unconditional.py | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index d6d8967b43ce..18040ac2f11d 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -272,7 +272,11 @@ def main(args): logging_dir=logging_dir, ) - if args.logger == "wandb": + if args.logger == "tensorboard": + if not is_tensorboard_available(): + raise ImportError("Make sure to install tensorboard if you want to use it for logging during training.") + + elif args.logger == "wandb": if not is_wandb_available(): raise ImportError("Make sure to install wandb if you want to use it for logging during training.") import wandb @@ -564,7 +568,7 @@ def transform_images(examples): # denormalize the images and save to tensorboard images_processed = (images * 255).round().astype("uint8") - if args.logger == "tensorboard" and is_tensorboard_available(): + if args.logger == "tensorboard": accelerator.get_tracker("tensorboard").add_images( "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index 8bdf6fc20dfd..b4789a14e493 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -280,7 +280,11 @@ def main(args): logging_dir=logging_dir, ) - if args.logger == "wandb": + if args.logger == "tensorboard": + if not is_tensorboard_available(): + raise ImportError("Make sure to install tensorboard if you want to use it for logging during training.") + + elif args.logger == "wandb": if not is_wandb_available(): raise ImportError("Make sure to install wandb if you want to use it for logging during training.") import wandb @@ -609,7 +613,7 @@ def transform_images(examples): # denormalize the images and save to tensorboard images_processed = (images * 255).round().astype("uint8") - if args.logger == "tensorboard" and is_tensorboard_available(): + if args.logger == "tensorboard": accelerator.get_tracker("tensorboard").add_images( "test_samples", images_processed.transpose(0, 3, 1, 2), epoch ) From 30b8ce73eb2899b7d58eab7bdeb7851f9780f6b7 Mon Sep 17 00:00:00 2001 From: Ben Evans Date: Tue, 14 Feb 2023 12:55:10 +0000 Subject: [PATCH 5/5] style fixes --- .../unconditional_image_generation/train_unconditional.py | 3 ++- examples/unconditional_image_generation/train_unconditional.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index 18040ac2f11d..cd120460c10f 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -574,7 +574,8 @@ def transform_images(examples): ) elif args.logger == "wandb": accelerator.get_tracker("wandb").log( - {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, step=global_step + {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, + step=global_step, ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index b4789a14e493..f76594a78c32 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -619,7 +619,8 @@ def transform_images(examples): ) elif args.logger == "wandb": accelerator.get_tracker("wandb").log( - {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, step=global_step + {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, + step=global_step, ) if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1: