Caffe2 - Python API
A deep learning, cross platform ML framework
relaxed_bernoulli.py
1 import torch
2 from numbers import Number
3 from torch.distributions import constraints
4 from torch.distributions.distribution import Distribution
5 from torch.distributions.transformed_distribution import TransformedDistribution
6 from torch.distributions.transforms import SigmoidTransform
7 from torch.distributions.utils import broadcast_all, probs_to_logits, logits_to_probs, lazy_property, clamp_probs
8 
9 
11  r"""
12  Creates a LogitRelaxedBernoulli distribution parameterized by :attr:`probs`
13  or :attr:`logits` (but not both), which is the logit of a RelaxedBernoulli
14  distribution.
15 
16  Samples are logits of values in (0, 1). See [1] for more details.
17 
18  Args:
19  temperature (Tensor): relaxation temperature
20  probs (Number, Tensor): the probability of sampling `1`
21  logits (Number, Tensor): the log-odds of sampling `1`
22 
23  [1] The Concrete Distribution: A Continuous Relaxation of Discrete Random
24  Variables (Maddison et al, 2017)
25 
26  [2] Categorical Reparametrization with Gumbel-Softmax
27  (Jang et al, 2017)
28  """
29  arg_constraints = {'probs': constraints.unit_interval,
30  'logits': constraints.real}
31  support = constraints.real
32 
33  def __init__(self, temperature, probs=None, logits=None, validate_args=None):
34  self.temperature = temperature
35  if (probs is None) == (logits is None):
36  raise ValueError("Either `probs` or `logits` must be specified, but not both.")
37  if probs is not None:
38  is_scalar = isinstance(probs, Number)
39  self.probs, = broadcast_all(probs)
40  else:
41  is_scalar = isinstance(logits, Number)
42  self.logits, = broadcast_all(logits)
43  self._param = self.probs if probs is not None else self.logits
44  if is_scalar:
45  batch_shape = torch.Size()
46  else:
47  batch_shape = self._param.size()
48  super(LogitRelaxedBernoulli, self).__init__(batch_shape, validate_args=validate_args)
49 
50  def expand(self, batch_shape, _instance=None):
51  new = self._get_checked_instance(LogitRelaxedBernoulli, _instance)
52  batch_shape = torch.Size(batch_shape)
53  new.temperature = self.temperature
54  if 'probs' in self.__dict__:
55  new.probs = self.probs.expand(batch_shape)
56  new._param = new.probs
57  else:
58  new.logits = self.logits.expand(batch_shape)
59  new._param = new.logits
60  super(LogitRelaxedBernoulli, new).__init__(batch_shape, validate_args=False)
61  new._validate_args = self._validate_args
62  return new
63 
64  def _new(self, *args, **kwargs):
65  return self._param.new(*args, **kwargs)
66 
67  @lazy_property
68  def logits(self):
69  return probs_to_logits(self.probs, is_binary=True)
70 
71  @lazy_property
72  def probs(self):
73  return logits_to_probs(self.logits, is_binary=True)
74 
75  @property
76  def param_shape(self):
77  return self._param.size()
78 
79  def rsample(self, sample_shape=torch.Size()):
80  shape = self._extended_shape(sample_shape)
81  probs = clamp_probs(self.probs.expand(shape))
82  uniforms = clamp_probs(torch.rand(shape, dtype=probs.dtype, device=probs.device))
83  return (uniforms.log() - (-uniforms).log1p() + probs.log() - (-probs).log1p()) / self.temperature
84 
85  def log_prob(self, value):
86  if self._validate_args:
87  self._validate_sample(value)
88  logits, value = broadcast_all(self.logits, value)
89  diff = logits - value.mul(self.temperature)
90  return self.temperature.log() + diff - 2 * diff.exp().log1p()
91 
92 
94  r"""
95  Creates a RelaxedBernoulli distribution, parametrized by
96  :attr:`temperature`, and either :attr:`probs` or :attr:`logits`
97  (but not both). This is a relaxed version of the `Bernoulli` distribution,
98  so the values are in (0, 1), and has reparametrizable samples.
99 
100  Example::
101 
102  >>> m = RelaxedBernoulli(torch.tensor([2.2]),
103  torch.tensor([0.1, 0.2, 0.3, 0.99]))
104  >>> m.sample()
105  tensor([ 0.2951, 0.3442, 0.8918, 0.9021])
106 
107  Args:
108  temperature (Tensor): relaxation temperature
109  probs (Number, Tensor): the probability of sampling `1`
110  logits (Number, Tensor): the log-odds of sampling `1`
111  """
112  arg_constraints = {'probs': constraints.unit_interval,
113  'logits': constraints.real}
114  support = constraints.unit_interval
115  has_rsample = True
116 
117  def __init__(self, temperature, probs=None, logits=None, validate_args=None):
118  base_dist = LogitRelaxedBernoulli(temperature, probs, logits)
119  super(RelaxedBernoulli, self).__init__(base_dist,
121  validate_args=validate_args)
122 
123  def expand(self, batch_shape, _instance=None):
124  new = self._get_checked_instance(RelaxedBernoulli, _instance)
125  return super(RelaxedBernoulli, self).expand(batch_shape, _instance=new)
126 
127  @property
128  def temperature(self):
129  return self.base_dist.temperature
130 
131  @property
132  def logits(self):
133  return self.base_dist.logits
134 
135  @property
136  def probs(self):
137  return self.base_dist.probs
def _get_checked_instance(self, cls, _instance=None)
def _extended_shape(self, sample_shape=torch.Size())