Upload layers.py

pgm/layers.py

@@ -149,27 +149,57 @@ class ArgMaxGumbelMax(Transform):
         """Infer the gumbels noises given k and logits."""
         assert self.logits != None, "Logits not defined."
 
-        uniforms = torch.rand(
-            self.logits.shape, dtype=self.logits.dtype, device=self.logits.device
-        )
-        gumbels = -((-(uniforms.log())).log())
-        # print(f'gumbels: {gumbels.size()}, {gumbels.dtype}')
-        # (batch_size, num_classes) mask to select kth class
-        # print(f'k : {k.size()}')
-        mask = F.one_hot(
-            k.squeeze(-1).to(torch.int64), num_classes=self.logits.shape[-1]
-        )
-        # print(f'mask: {mask.size()}, {mask.dtype}')
-        # (batch_size, 1) select topgumbel for truncation of other classes
-        topgumbel = (mask * gumbels).sum(dim=-1, keepdim=True) - (
-            mask * self.logits
-        ).sum(dim=-1, keepdim=True)
-        mask = 1 - mask  # invert mask to select other != k classes
-        g = gumbels + self.logits
-        # (batch_size, num_classes)
-        epsilons = -torch.log(mask * torch.exp(-g) + torch.exp(-topgumbel)) - (
-            mask * self.logits
-        )
+        # uniforms = torch.rand(
+        #     self.logits.shape, dtype=self.logits.dtype, device=self.logits.device
+        # )
+        # gumbels = -((-(uniforms.log())).log())
+        # # print(f'gumbels: {gumbels.size()}, {gumbels.dtype}')
+        # # (batch_size, num_classes) mask to select kth class
+        # # print(f'k : {k.size()}')
+        # mask = F.one_hot(
+        #     k.squeeze(-1).to(torch.int64), num_classes=self.logits.shape[-1]
+        # )
+        # # print(f'mask: {mask.size()}, {mask.dtype}')
+        # # (batch_size, 1) select topgumbel for truncation of other classes
+        # topgumbel = (mask * gumbels).sum(dim=-1, keepdim=True) - (
+        #     mask * self.logits
+        # ).sum(dim=-1, keepdim=True)
+        # mask = 1 - mask  # invert mask to select other != k classes
+        # g = gumbels + self.logits
+        # # (batch_size, num_classes)
+        # epsilons = -torch.log(mask * torch.exp(-g) + torch.exp(-topgumbel)) - (
+        #     mask * self.logits
+        # )
+        
+        def sample_gumbel(shape, eps=1e-20):
+            U = torch.rand(shape)
+            
+            U = U.cuda()
+            return -torch.log(-torch.log(U + eps) + eps)
+        def gumbel_softmax_sample(logits, temperature):
+            y = logits + sample_gumbel(logits.shape)
+            return F.softmax(y / temperature, dim=-1)
+        def gumbel_softmax(logits, temperature,k, hard=False):
+            """
+            ST-gumple-softmax
+            input: [*, n_class]
+            return: flatten --> [*, n_class] an one-hot vector
+            """
+            y = gumbel_softmax_sample(logits, temperature)
+        
+            if not hard:
+                return y.view(-1, logits.shape[-1])
+        
+            shape = y.size()
+            _, ind = k.max(dim=-1)
+            y_hard = torch.zeros_like(y).view(-1, shape[-1])
+            y_hard.scatter_(1, ind.view(-1, 1), 1)
+            y_hard = y_hard.view(*shape)
+            # Set gradients w.r.t. y_hard gradients w.r.t. y
+            y_hard = (y_hard - y).detach() + y
+            return y_hard.view(-1, logits.shape[-1])
+        epsilons = gumbel_softmax(self.logits,1e-3,k)
+        
         return epsilons
 
     def log_abs_det_jacobian(self, x, y):