Made changes in the training script

celebhq/vqvae_autoencoder_ckpt.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cb6f542f508f916324196d1d64aa13657905de80419b8b638387bd67040f0ed5
+oid sha256:34e9a28cd01122493eed58be315fb4757b3513701b47a4bf6058a59506f3f7fd
 size 88110787

celebhq/vqvae_optim_g_ckpt.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6f8ae23a0f51866d3116b110a66b983f271e6461c5892c1d945870d84a1fcee4
+oid sha256:f29a6752be23312bbe39525029bd015a401b32d25de7b38f3200b1e0473a22ef
 size 264209698

config/celebahq.yaml

@@ -49,7 +49,7 @@ train_config:
   ldm_batch_size: 16
   autoencoder_batch_size: 4
   disc_start: 15000
-  disc_weight: 0.5
+  disc_beta: 0.5
   codebook_beta: 1
   commitment_beta: 0.2
   perceptual_weight: 1

train_vqvae.py

@@ -53,8 +53,8 @@ def train(args):
         print("Loading checkpoint...")
         model = torch.load(vqvae_ckpt_path).to(device)
         discriminator = torch.load(discriminator_ckpt_path).to(device)
-        optimizer_d = torch.load(optimizer_d_ckpt).to(device)
-        optimizer_g = torch.load(optimizer_g_ckpt).to(device)
+        optimizer_d = torch.load(optimizer_d_ckpt)
+        optimizer_g = torch.load(optimizer_g_ckpt)
         
     else:
         model = VQVAE(
@@ -69,10 +69,10 @@ def train(args):
     
 
     img_save_steps = train_config["autoencoder_img_save_steps"]
-    img_saved = 0
+    img_saved = 14
 
     disc_step_start = train_config['disc_start']
-    steps = 0
+    steps = 15000
 
     for epoch in range(train_config["autoencoder_epochs"]):
         recon_losses = []
@@ -129,7 +129,7 @@ def train(args):
             if steps > disc_step_start:
                 disc_fake_pred = discriminator(model_output[0])
                 disc_fake_loss = disc_criterion(
-                    disc_fake_pred, torch.ones_like(disc_fake_pred.shape, device=disc_fake_pred.device))
+                    disc_fake_pred, torch.ones(disc_fake_pred.shape, device=disc_fake_pred.device))
                 gen_losses.append(
                     train_config["disc_beta"] * disc_fake_loss.item())
                 g_loss += train_config["disc_beta"] * disc_fake_loss
@@ -149,7 +149,7 @@ def train(args):
                     disc_fake_pred, device=disc_fake_pred.device))
                 disc_real_loss = disc_criterion(disc_real_pred, torch.ones_like(
                     disc_real_pred, device=disc_real_pred.device))
-                disc_loss = train_config["disc_weight"] * \
+                disc_loss = train_config["disc_beta"] * \
                     (disc_real_loss + disc_fake_loss) / 2
                 disc_losses.append(disc_loss)
                 disc_loss.backward()