Categorical random variable can use gumbel-softmax sampling
@@ -6,6 +6,7 @@ from __future__ import unicode_literals
import torch
from torch.autograd import Variable
from .random_variable import RandomVariable
+from torch.nn import functional as F
class CategoricalRandomVariable(RandomVariable):
@@ -45,31 +46,38 @@ class CategoricalRandomVariable(RandomVariable):
def representation(self):
return self.mass_function
- def sample(self, n_samples=1):
+ def sample(self, n_samples=1, gumbel_softmax_temp=None):
"""
This function is written as a reparameterization
This way we can differentiate through the mass function
"""
mass_function = self.mass_function
+
if mass_function.ndimension() == 1:
mass_function = mass_function.unsqueeze(0)
+ # Generate uniform samples
+
+ if gumbel_softmax_temp is None:
upper_mass_function = mass_function.cumsum(1)
lower_mass_function = upper_mass_function - mass_function
- # Generate uniform samples
- samples = Variable(mass_function.data.new(n_samples, 1, 1).uniform_())
+ samples = Variable(mass_function.data.new(1, 1, n_samples).uniform_())
samples = samples.clamp(1e-5, 1) # Make sure that everything is strictly greater than zero
-
- lower_mask = samples.gt(lower_mass_function.unsqueeze(0))
- upper_mask = samples.le(upper_mass_function.unsqueeze(0))
+ lower_mask = samples.gt(lower_mass_function.unsqueeze(-1))
+ upper_mask = samples.le(upper_mass_function.unsqueeze(-1))
res = (lower_mask * upper_mask)
+ else:
+ samples = Variable(mass_function.data.new(1, mass_function.size(-1), n_samples).uniform_())
+ gumbel_samples = samples.log().mul_(-1).log_().mul_(-1)
+ softmax_weights = (mass_function.log().unsqueeze(-1) + gumbel_samples).div(gumbel_softmax_temp)
+ res = F.softmax(softmax_weights, 1)
# Sample dimension is first
if self.mass_function.ndimension() == 1:
- res = res.squeeze(1)
- if n_samples == 1:
res = res.squeeze(0)
+ if n_samples == 1:
+ res = res.squeeze(-1)
return res
def __len__(self):
