Log Unconditional Image Generation Samples to WandB

Author: bencevans

examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py

@@ -561,6 +561,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

examples/unconditional_image_generation/train_unconditional.py

@@ -596,6 +596,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