Caffe2 - Python API A deep learning, cross platform ML framework
categorical.py
1 import torch
2 from torch._six import nan
3 from torch.distributions import constraints
4 from torch.distributions.distribution import Distribution
5 from torch.distributions.utils import probs_to_logits, logits_to_probs, lazy_property, broadcast_all
6
7
9  r"""
10  Creates a categorical distribution parameterized by either :attr:`probs` or
11  :attr:`logits` (but not both).
12
13  .. note::
14  It is equivalent to the distribution that :func:`torch.multinomial`
15  samples from.
16
17  Samples are integers from :math:`\{0, \ldots, K-1\}` where `K` is ``probs.size(-1)``.
18
19  If :attr:`probs` is 1D with length-`K`, each element is the relative
20  probability of sampling the class at that index.
21
22  If :attr:`probs` is 2D, it is treated as a batch of relative probability
23  vectors.
24
25  .. note:: :attr:`probs` must be non-negative, finite and have a non-zero sum,
26  and it will be normalized to sum to 1.
27
28  See also: :func:`torch.multinomial`
29
30  Example::
31
32  >>> m = Categorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ]))
33  >>> m.sample() # equal probability of 0, 1, 2, 3
34  tensor(3)
35
36  Args:
37  probs (Tensor): event probabilities
38  logits (Tensor): event log probabilities
39  """
40  arg_constraints = {'probs': constraints.simplex,
41  'logits': constraints.real}
42  has_enumerate_support = True
43
44  def __init__(self, probs=None, logits=None, validate_args=None):
45  if (probs is None) == (logits is None):
46  raise ValueError("Either `probs` or `logits` must be specified, but not both.")
47  if probs is not None:
48  if probs.dim() < 1:
49  raise ValueError("`probs` parameter must be at least one-dimensional.")
50  self.probs = probs / probs.sum(-1, keepdim=True)
51  else:
52  if logits.dim() < 1:
53  raise ValueError("`logits` parameter must be at least one-dimensional.")
54  self.logits = logits - logits.logsumexp(dim=-1, keepdim=True)
55  self._param = self.probs if probs is not None else self.logits
56  self._num_events = self._param.size()[-1]
57  batch_shape = self._param.size()[:-1] if self._param.ndimension() > 1 else torch.Size()
58  super(Categorical, self).__init__(batch_shape, validate_args=validate_args)
59
60  def expand(self, batch_shape, _instance=None):
61  new = self._get_checked_instance(Categorical, _instance)
62  batch_shape = torch.Size(batch_shape)
63  param_shape = batch_shape + torch.Size((self._num_events,))
64  if 'probs' in self.__dict__:
65  new.probs = self.probs.expand(param_shape)
66  new._param = new.probs
67  else:
68  new.logits = self.logits.expand(param_shape)
69  new._param = new.logits
70  new._num_events = self._num_events
71  super(Categorical, new).__init__(batch_shape, validate_args=False)
72  new._validate_args = self._validate_args
73  return new
74
75  def _new(self, *args, **kwargs):
76  return self._param.new(*args, **kwargs)
77
78  @constraints.dependent_property
79  def support(self):
80  return constraints.integer_interval(0, self._num_events - 1)
81
82  @lazy_property
83  def logits(self):
84  return probs_to_logits(self.probs)
85
86  @lazy_property
87  def probs(self):
88  return logits_to_probs(self.logits)
89
90  @property
91  def param_shape(self):
92  return self._param.size()
93
94  @property
95  def mean(self):
96  return self.probs.new_tensor(nan).expand(self._extended_shape())
97
98  @property
99  def variance(self):
100  return self.probs.new_tensor(nan).expand(self._extended_shape())
101
102  def sample(self, sample_shape=torch.Size()):
103  sample_shape = self._extended_shape(sample_shape)
104  param_shape = sample_shape + torch.Size((self._num_events,))
105  probs = self.probs.expand(param_shape)
106  if self.probs.dim() == 1 or self.probs.size(0) == 1:
107  probs_2d = probs.view(-1, self._num_events)
108  else:
109  probs_2d = probs.contiguous().view(-1, self._num_events)
110  sample_2d = torch.multinomial(probs_2d, 1, True)
111  return sample_2d.contiguous().view(sample_shape)
112
113  def log_prob(self, value):
114  if self._validate_args:
115  self._validate_sample(value)
116  value = value.long().unsqueeze(-1)
117  value, log_pmf = torch.broadcast_tensors(value, self.logits)
118  value = value[..., :1]
119  return log_pmf.gather(-1, value).squeeze(-1)
120
121  def entropy(self):
122  p_log_p = self.logits * self.probs
123  return -p_log_p.sum(-1)
124
125  def enumerate_support(self, expand=True):
126  num_events = self._num_events
127  values = torch.arange(num_events, dtype=torch.long, device=self._param.device)
128  values = values.view((-1,) + (1,) * len(self._batch_shape))
129  if expand:
130  values = values.expand((-1,) + self._batch_shape)
131  return values
def _get_checked_instance(self, cls, _instance=None)
def _extended_shape(self, sample_shape=torch.Size())