kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix test train_unconditional (#2481)

Browse Source 
* Fix tensorboard tracking with `accelerate` @ `main`

* Fix `train_unconditional.py` with accelerate from main.
This commit is contained in:
Pedro Cuenca
2023-02-24 23:31:16 +01:00
committed by GitHub
parent 54bc882d96
commit 5de4347663
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/unconditional_image_generation/train_unconditional.py
+7 -4
View File
@@ -23,7 +23,7 @@ import diffusers
from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel
from diffusers.optimization import get_scheduler from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel from diffusers.training_utils import EMAModel
from diffusers.utils import check_min_version, is_tensorboard_available, is_wandb_available from diffusers.utils import check_min_version, is_accelerate_version, is_tensorboard_available, is_wandb_available
# Will error if the minimal version of diffusers is not installed. Remove at your own risks. # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
@@ -628,10 +628,13 @@ def main(args):
images_processed = (images * 255).round().astype("uint8") images_processed = (images * 255).round().astype("uint8")
if args.logger == "tensorboard": if args.logger == "tensorboard":
accelerator.get_tracker("tensorboard").add_images( if is_accelerate_version(">=", "0.17.0.dev0"):
"test_samples", images_processed.transpose(0, 3, 1, 2), epoch tracker = accelerator.get_tracker("tensorboard", unwrap=True)
) else:
tracker = accelerator.get_tracker()
tracker.add_images("test_samples", images_processed.transpose(0, 3, 1, 2), epoch)
elif args.logger == "wandb": elif args.logger == "wandb":
# Upcoming `log_images` helper coming in https://github.com/huggingface/accelerate/pull/962/files
accelerator.get_tracker("wandb").log( accelerator.get_tracker("wandb").log(
{"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch}, {"test_samples": [wandb.Image(img) for img in images_processed], "epoch": epoch},
step=global_step, step=global_step,