Return indices in vq_st

Author: tristandeleu

functions.py

@@ -40,10 +40,10 @@ def forward(ctx, inputs, codebook):
             index=indices_flatten)
         codes = codes_flatten.view_as(inputs)
 
-        return codes
+        return (codes, indices_flatten)
 
     @staticmethod
-    def backward(ctx, grad_output):
+    def backward(ctx, grad_output, grad_indices):
         grad_inputs, grad_codebook = None, None
 
         if ctx.needs_input_grad[0]:

modules.py

@@ -80,9 +80,15 @@ def forward(self, z_e_x):
 
     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_, indices = vq_st(z_e_x_, self.embedding.weight)
         z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()
-        return z_q_x
+
+        z_q_x_bar_flatten = torch.index_select(self.embedding.weight,
+            dim=0, index=indices)
+        z_q_x_bar_ = z_q_x_bar_flatten.view_as(z_e_x_)
+        z_q_x_bar = z_q_x_bar_.permute(0, 3, 1, 2).contiguous()
+
+        return z_q_x, z_q_x_bar
 
 
 class ResBlock(nn.Module):
@@ -140,8 +146,8 @@ def decode(self, latents):
 
     def forward(self, x):
         z_e_x = self.encoder(x)
-        z_q_x = self.codebook.straight_through(z_e_x)
-        x_tilde = self.decoder(z_q_x)
+        z_q_x_st, z_q_x = self.codebook.straight_through(z_e_x)
+        x_tilde = self.decoder(z_q_x_st)
         return x_tilde, z_e_x, z_q_x