gumbel_softmax tweaks

hughperkins · hughperkins · commit 53f43a769349 · 2017-12-29T22:55:00.000-05:00
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -1016,8 +1016,6 @@ def test_embedding_functional(self):
 
     def _test_gumbel_softmax_st(self, cuda):
         th = torch.cuda if cuda else torch
-        old_rng_state = th.get_rng_state()
-        th.manual_seed(42)
         """
         Things we might want to check:
         - if we make various draws, do we get different one-hot values?
@@ -1037,6 +1035,7 @@ def _test_gumbel_softmax_st(self, cuda):
             y_draws = y_draws.cuda()
             preds = preds.cuda()
 
+        exceed_limits = 0
         for draw in range(num_draws):
             logits_var = Variable(logits, requires_grad=True)
             y_draw = torch.nn.functional.gumbel_softmax(
@@ -1048,10 +1047,12 @@ def _test_gumbel_softmax_st(self, cuda):
             err = y_draw - Variable(logits.new([[0, 0.5, 0.3]]))
             loss = (err * err).sum()
             loss.backward()
-            assert logits_var.grad.abs().min().data[0] > 0.001
+            if logits_var.grad.data.std() < 0.01 or logits_var.grad.data.std() > 1.0:
+                exceed_limits += 1
             y_draws[draw] = y_draw.data
             _, pred = y_draw.max(1)
             preds[draw] = pred.data[0]
+        assert exceed_limits / num_draws < 0.05
         # check it's approximately one-hot
         num_ones = (y_draws == 1).int().sum()
         num_zeros = (y_draws == 0).int().sum()
@@ -1061,7 +1062,6 @@ def _test_gumbel_softmax_st(self, cuda):
         num_class_one = (preds == 1).int().sum()
         assert num_class_one < num_draws
         assert num_class_one > num_draws / 3
-        th.set_rng_state(old_rng_state)
 
     def test_gumbel_softmax_st(self):
         self._test_gumbel_softmax_st(False)
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -835,7 +835,7 @@ def softmax(input, dim=None, _stacklevel=3):
     return torch._C._nn.softmax(input, dim)
 
 
-def sample_gumbel(shape, eps=1e-10, out=None):
+def _sample_gumbel(shape, eps=1e-10, out=None):
     """
     Sample from Gumbel(0, 1)
 
@@ -847,16 +847,16 @@ def sample_gumbel(shape, eps=1e-10, out=None):
     return - torch.log(eps - torch.log(U + eps))
 
 
-def gumbel_softmax_sample(logits, tau=1, eps=1e-10):
+def _gumbel_softmax_sample(logits, tau=1, eps=1e-10):
     """
     Draw a sample from the Gumbel-Softmax distribution
 
     based on
     https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb
     (MIT license)
     """
-    dims = len(logits.size())
-    gumbel_noise = sample_gumbel(logits.size(), eps=eps, out=logits.data.new())
+    dims = logits.dim()
+    gumbel_noise = _sample_gumbel(logits.size(), eps=eps, out=logits.data.new())
     y = logits + Variable(gumbel_noise)
     return softmax(y / tau, dims - 1)
 
@@ -882,7 +882,7 @@ def gumbel_softmax(logits, tau=1, hard=False, eps=1e-10):
     """
     shape = logits.size()
     assert len(shape) == 2
-    y_soft = gumbel_softmax_sample(logits, tau=tau, eps=eps)
+    y_soft = _gumbel_softmax_sample(logits, tau=tau, eps=eps)
     if hard:
         _, k = y_soft.data.max(-1)
         # this bit is based on
