Skip to content

Commit 753cb01

Browse files
committed
Log Unconditional Image Generation Samples to WandB
1 parent 9d0d070 commit 753cb01

File tree

2 files changed

+10
-0
lines changed

2 files changed

+10
-0
lines changed

examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,11 @@ def transform_images(examples):
561561
accelerator.get_tracker("tensorboard").add_images(
562562
"test_samples", images_processed.transpose(0, 3, 1, 2), epoch
563563
)
564+
elif args.logger == "wandb":
565+
import wandb
566+
accelerator.get_tracker("wandb").log(
567+
{"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step
568+
)
564569

565570
if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
566571
# save the model

examples/unconditional_image_generation/train_unconditional.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -596,6 +596,11 @@ def transform_images(examples):
596596
accelerator.get_tracker("tensorboard").add_images(
597597
"test_samples", images_processed.transpose(0, 3, 1, 2), epoch
598598
)
599+
elif args.logger == "wandb":
600+
import wandb
601+
accelerator.get_tracker("wandb").log(
602+
{"test_samples": [wandb.Image(img) for img in images_processed]}, step=global_step
603+
)
599604

600605
if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
601606
# save the model

0 commit comments

Comments
 (0)