def reparmeterize(self, logits):
""" reparamterize the logits
:param logits: non-activated logits
:returns: reparameterized and hard outputs
:rtype: torch.Tensor, torch.Tensor
"""
relaxed = D.RelaxedBernoulli(temperature=self.tau,
logits=logits).rsample()
hard = self._compute_hard(relaxed)
return relaxed, hard
def reparmeterize(self, logits):
""" reparamterize the logits
:param logits: non-activated logits
:returns: reparameterized and hard outputs
:rtype: torch.Tensor, torch.Tensor
"""
relaxed = D.RelaxedBernoulli(temperature=self.tau, logits=logits).rsample()
hard = relaxed.clone()
hard[relaxed < 0.5] = 0.0
hard[relaxed >= 0.5] = 1.0
return relaxed, hard
def _build_posterior(self, param_logit):
dist_posterior = distributions.RelaxedBernoulli(
logits=torch.clamp_min(param_logit, _LOG_EPSILON),
temperature=self.temperature)
return dist_posterior
def _build_prior(self, prior_logit, prior_temperature):
dist_prior = distributions.RelaxedBernoulli(
logits=torch.clamp_min(prior_logit, _LOG_EPSILON),
temperature=torch.tensor(prior_temperature, dtype=torch.float32))
return dist_prior
def sample(self, batch_size):
model = dists.RelaxedBernoulli(self.temperature, self.probs)
return model.sample((batch_size,))
def log_prob(self, value):
model = dists.RelaxedBernoulli(self.temperature, self.probs)
return model.log_prob(value).sum(-1)