pytorch

1
import torch
2
from torch.distributions import constraints
3
from torch.distributions.categorical import Categorical
4
from torch.distributions.distribution import Distribution
5

6
__all__ = ["OneHotCategorical", "OneHotCategoricalStraightThrough"]
7

8

9
class OneHotCategorical(Distribution):
10
    r"""
11
    Creates a one-hot categorical distribution parameterized by :attr:`probs` or
12
    :attr:`logits`.
13

14
    Samples are one-hot coded vectors of size ``probs.size(-1)``.
15

16
    .. note:: The `probs` argument must be non-negative, finite and have a non-zero sum,
17
              and it will be normalized to sum to 1 along the last dimension. :attr:`probs`
18
              will return this normalized value.
19
              The `logits` argument will be interpreted as unnormalized log probabilities
20
              and can therefore be any real number. It will likewise be normalized so that
21
              the resulting probabilities sum to 1 along the last dimension. :attr:`logits`
22
              will return this normalized value.
23

24
    See also: :func:`torch.distributions.Categorical` for specifications of
25
    :attr:`probs` and :attr:`logits`.
26

27
    Example::
28

29
        >>> # xdoctest: +IGNORE_WANT("non-deterministic")
30
        >>> m = OneHotCategorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ]))
31
        >>> m.sample()  # equal probability of 0, 1, 2, 3
32
        tensor([ 0.,  0.,  0.,  1.])
33

34
    Args:
35
        probs (Tensor): event probabilities
36
        logits (Tensor): event log probabilities (unnormalized)
37
    """
38
    arg_constraints = {"probs": constraints.simplex, "logits": constraints.real_vector}
39
    support = constraints.one_hot
40
    has_enumerate_support = True
41

42
    def __init__(self, probs=None, logits=None, validate_args=None):
43
        self._categorical = Categorical(probs, logits)
44
        batch_shape = self._categorical.batch_shape
45
        event_shape = self._categorical.param_shape[-1:]
46
        super().__init__(batch_shape, event_shape, validate_args=validate_args)
47

48
    def expand(self, batch_shape, _instance=None):
49
        new = self._get_checked_instance(OneHotCategorical, _instance)
50
        batch_shape = torch.Size(batch_shape)
51
        new._categorical = self._categorical.expand(batch_shape)
52
        super(OneHotCategorical, new).__init__(
53
            batch_shape, self.event_shape, validate_args=False
54
        )
55
        new._validate_args = self._validate_args
56
        return new
57

58
    def _new(self, *args, **kwargs):
59
        return self._categorical._new(*args, **kwargs)
60

61
    @property
62
    def _param(self):
63
        return self._categorical._param
64

65
    @property
66
    def probs(self):
67
        return self._categorical.probs
68

69
    @property
70
    def logits(self):
71
        return self._categorical.logits
72

73
    @property
74
    def mean(self):
75
        return self._categorical.probs
76

77
    @property
78
    def mode(self):
79
        probs = self._categorical.probs
80
        mode = probs.argmax(axis=-1)
81
        return torch.nn.functional.one_hot(mode, num_classes=probs.shape[-1]).to(probs)
82

83
    @property
84
    def variance(self):
85
        return self._categorical.probs * (1 - self._categorical.probs)
86

87
    @property
88
    def param_shape(self):
89
        return self._categorical.param_shape
90

91
    def sample(self, sample_shape=torch.Size()):
92
        sample_shape = torch.Size(sample_shape)
93
        probs = self._categorical.probs
94
        num_events = self._categorical._num_events
95
        indices = self._categorical.sample(sample_shape)
96
        return torch.nn.functional.one_hot(indices, num_events).to(probs)
97

98
    def log_prob(self, value):
99
        if self._validate_args:
100
            self._validate_sample(value)
101
        indices = value.max(-1)[1]
102
        return self._categorical.log_prob(indices)
103

104
    def entropy(self):
105
        return self._categorical.entropy()
106

107
    def enumerate_support(self, expand=True):
108
        n = self.event_shape[0]
109
        values = torch.eye(n, dtype=self._param.dtype, device=self._param.device)
110
        values = values.view((n,) + (1,) * len(self.batch_shape) + (n,))
111
        if expand:
112
            values = values.expand((n,) + self.batch_shape + (n,))
113
        return values
114

115

116
class OneHotCategoricalStraightThrough(OneHotCategorical):
117
    r"""
118
    Creates a reparameterizable :class:`OneHotCategorical` distribution based on the straight-
119
    through gradient estimator from [1].
120

121
    [1] Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
122
    (Bengio et al, 2013)
123
    """
124
    has_rsample = True
125

126
    def rsample(self, sample_shape=torch.Size()):
127
        samples = self.sample(sample_shape)
128
        probs = self._categorical.probs  # cached via @lazy_property
129
        return samples + (probs - probs.detach())
130