log images in tensorboard #998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

bghira merged 27 commits into bghira:main from anhi:release

Sep 28, 2024

helpers/training/validation.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -2,6 +2,7 @@ @@
     import os
     import wandb
     import logging
+    import numpy as np
     from tqdm import tqdm
     from helpers.training.wrappers import unwrap_model
     from PIL import Image
@@ Expand Down Expand Up @@
                                 image,
                                 name=f"{shortname} - {self.validation_resolutions[idx]}",
                             )
+                elif tracker.name == "tensorboard":
+                    tracker = self.accelerator.get_tracker("tensorboard")
+                    for shortname, image_list in validation_images.items():
+                        tracker.log_images(
+                                {
+                                    f"{shortname} - {self.validation_resolutions[idx]}": np.moveaxis(np.array(image), -1, 0)[np.newaxis, ...]
+                                    for idx, image in enumerate(image_list)
+                                },
+                                step=StateTracker.get_global_step()
+                        )
                 elif tracker.name == "wandb":
                     resolution_list = [
                         f"{res[0]}x{res[1]}" for res in get_validation_resolutions()
@@ Expand Down @@

Provide feedback