Replace quantization in VQEmbedding by vq function

Author: tristandeleu

miniimagenet_pixelcnn_prior.py

@@ -14,7 +14,7 @@ def train(data_loader, model, prior, optimizer, args, writer):
     for images, labels in data_loader:
         with torch.no_grad():
             images = images.to(args.device)
-            latents, _ = model.encode(images)
+            latents = model.encode(images)
             latents = latents.detach()
 
         labels = labels.to(args.device)
@@ -39,7 +39,7 @@ def test(data_loader, model, prior, args, writer):
             images = images.to(args.device)
             labels = labels.to(args.device)
 
-            latents, _ = model.encode(images)
+            latents = model.encode(images)
             latents = latents.detach()
             logits = prior(latents, labels)
             logits = logits.permute(0, 2, 3, 1).contiguous()

miniimagenet_vqvae.py

@@ -15,22 +15,16 @@ def train(data_loader, model, optimizer, args, writer):
 
         optimizer.zero_grad()
         x_tilde, z_e_x, z_q_x = model(images)
-        z_q_x.retain_grad()
 
+        # Reconstruction loss
         loss_recons = F.mse_loss(x_tilde, images)
-        loss_recons.backward(retain_graph=True)
-
-        # Straight-through estimator
-        z_e_x.backward(z_q_x.grad, retain_graph=True)
-
         # Vector quantization objective
-        model.codebook.embedding.zero_grad()
         loss_vq = F.mse_loss(z_q_x, z_e_x.detach())
-        loss_vq.backward(retain_graph=True)
-
         # Commitment objective
-        loss_commit = args.beta * F.mse_loss(z_e_x, z_q_x.detach())
-        loss_commit.backward()
+        loss_commit = F.mse_loss(z_e_x, z_q_x.detach())
+
+        loss = loss_recons + loss_vq + args.beta * loss_commit
+        loss.backward()
 
         # Logs
         writer.add_scalar('loss/train/reconstruction', loss_recons.item(), args.steps)

modules.py

@@ -4,6 +4,7 @@
 from torch.distributions.normal import Normal
 from torch.distributions import kl_divergence
 
+from functions import vq, vq_st
 
 def to_scalar(arr):
     if type(arr) == list:
@@ -73,18 +74,16 @@ def __init__(self, K, D):
         self.embedding.weight.data.uniform_(-1./K, 1./K)
 
     def forward(self, z_e_x):
-        # z_e_x - (B, D, H, W)
-        # emb   - (K, D)
-
-        emb = self.embedding.weight
-        dists = torch.pow(
-            z_e_x.unsqueeze(1) - emb[None, :, :, None, None],
-            2
-        ).sum(2)
-
-        latents = dists.min(1)[1]
+        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
+        latents = vq(z_e_x_, self.embedding.weight)
         return latents
 
+    def straight_through(self, z_e_x):
+        z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
+        z_q_x_ = vq_st(z_e_x_, self.embedding.weight)
+        z_q_x = z_q_x_.permute(0, 3, 1, 2)
+        return z_q_x
+
 
 class ResBlock(nn.Module):
     def __init__(self, dim):
@@ -132,16 +131,17 @@ def __init__(self, input_dim, dim, K=512):
     def encode(self, x):
         z_e_x = self.encoder(x)
         latents = self.codebook(z_e_x)
-        return latents, z_e_x
+        return latents
 
     def decode(self, latents):
         z_q_x = self.codebook.embedding(latents).permute(0, 3, 1, 2)  # (B, D, H, W)
         x_tilde = self.decoder(z_q_x)
-        return x_tilde, z_q_x
+        return x_tilde
 
     def forward(self, x):
-        latents, z_e_x = self.encode(x)
-        x_tilde, z_q_x = self.decode(latents)
+        z_e_x = self.encoder(x)
+        z_q_x = self.codebook.straight_through(z_e_x)
+        x_tilde = self.decoder(z_q_x)
         return x_tilde, z_e_x, z_q_x
 
 

pixelcnn_prior.py

@@ -118,7 +118,7 @@ def generate_samples():
     label = label.to(device=DEVICE, dtype=torch.int64)
 
     latents = model.generate(label, shape=LATENT_SHAPE, batch_size=100)
-    x_tilde, _ = autoencoder.decode(latents)
+    x_tilde = autoencoder.decode(latents)
     images = (x_tilde.cpu().data + 1) / 2
 
     save_image(