Add wandb logging (Life is more peacefull)

train_vqvae.py

@@ -2,17 +2,19 @@ import os
 import argparse
 import torch
 import torch.nn as nn
-import torchvision
-from torchvision import models
 from models.vqvae import VQVAE
 from models.discriminator import Discriminator
 from torch.optim import Adam
 from models.lpips import LPIPS
 from dataset.celeba import create_dataloader
 from torchvision.utils import make_grid
+from torchvision.transforms import ToPILImage
 import yaml
 import numpy as np
 from tqdm import tqdm
+import wandb
+
+wandb.init(project="vqvae")
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -77,6 +79,8 @@ def train(args):
         optimizer_g = Adam(
             model.parameters(), lr=train_config["autoencoder_lr"], betas=(0.5, 0.999)
         )
+    wandb.watch(model,log="all", log_freq=100)
+
     # LPIPS model
     lpips_model = LPIPS().eval().to(device)
 
@@ -115,24 +119,10 @@ def train(args):
                 grid = make_grid(
                     torch.cat([save_input, save_output], dim=0), nrow=sample_size
                 )
-                img = torchvision.transforms.ToPILImage()(grid)
-                if not os.path.exists(
-                    os.path.join(train_config["task_name"], "vqvae_autoencoder_samples")
-                ):
-                    os.mkdir(
-                        os.path.join(
-                            train_config["task_name"], "vqvae_autoencoder_samples"
-                        )
-                    )
-                img.save(
-                    os.path.join(
-                        train_config["task_name"],
-                        "vqvae_autoencoder_samples",
-                        "current_autoencoder_sample_{}.png".format(img_saved),
-                    )
-                )
+                
+                grid_image = ToPILImage(grid)
+                wandb.log({"Latent generation": wandb.Image(grid_image,caption=f"Epoch: {epoch+1}, Step: {steps}")})
                 img_saved += 1
-                img.close()
 
             steps += 1
 
@@ -189,6 +179,19 @@ def train(args):
         optimizer_g.zero_grad()
         optimizer_d.step()
         optimizer_d.zero_grad()
+
+        wandb.log({
+            "epoch": epoch + 1,
+            "step": steps,
+            "image_saved": img_saved,
+            "recon_loss": np.mean(recon_losses),
+            "perceptual_loss": np.mean(perceptual_losses),
+            "codebook_loss": np.mean(codebook_losses),
+            "gen_loss": np.mean(gen_losses),
+            "disc_loss": np.mean(disc_losses),
+            "overall_loss": np.mean(losses)
+        })
+
         if len(disc_losses) > 0:
             print(
                 "Finished epoch: {} | Recon Loss : {:.4f} | Perceptual Loss : {:.4f} | "
@@ -235,7 +238,29 @@ def train(args):
             ),
         )
 
+        wandb.save(
+            os.path.join(
+                train_config["task_name"], train_config["vqvae_autoencoder_ckpt_name"]
+            )
+        )
+        wandb.save(
+            os.path.join(
+                train_config["task_name"], train_config["vqvae_discriminator_ckpt_name"]
+            )
+        )
+        wandb.save(
+            os.path.join(
+                train_config["task_name"], train_config["vqvae_optim_d_ckpt_name"]
+            )
+        )
+        wandb.save(
+            os.path.join(
+                train_config["task_name"], train_config["vqvae_optim_g_ckpt_name"]
+            )
+        )
+
     print("Done Training....")
+    wandb.finish()
 
 
 if __name__ == "__main__":