507 lines
16 KiB
Python
507 lines
16 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle.base.data_feeder import check_type, convert_dtype
|
|
from paddle.base.framework import Variable
|
|
from paddle.distribution import exponential_family
|
|
from paddle.framework import in_dynamic_mode
|
|
from paddle.nn.functional import (
|
|
binary_cross_entropy_with_logits,
|
|
sigmoid,
|
|
softplus,
|
|
)
|
|
from paddle.utils.decorator_utils import param_one_alias
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Sequence
|
|
|
|
from paddle import Tensor
|
|
from paddle._typing.dtype_like import _DTypeLiteral
|
|
|
|
# Smallest representable number
|
|
EPS = {
|
|
'float32': paddle.finfo(paddle.float32).eps,
|
|
'float64': paddle.finfo(paddle.float64).eps,
|
|
}
|
|
|
|
|
|
def _clip_probs(probs, dtype):
|
|
"""Clip probs from [0, 1] to (0, 1) with ``eps``.
|
|
|
|
Args:
|
|
probs (Tensor): probs of Bernoulli.
|
|
dtype (str): data type.
|
|
|
|
Returns:
|
|
Tensor: Clipped probs.
|
|
"""
|
|
eps = EPS.get(dtype)
|
|
return paddle.clip(probs, min=eps, max=1 - eps).astype(dtype)
|
|
|
|
|
|
class Bernoulli(exponential_family.ExponentialFamily):
|
|
r"""Bernoulli distribution parameterized by ``probs``, which is the probability of value 1.
|
|
|
|
In probability theory and statistics, the Bernoulli distribution, named after Swiss
|
|
mathematician Jacob Bernoulli, is the discrete probability distribution of a random
|
|
variable which takes the value 1 with probability ``p`` and the value 0 with
|
|
probability ``q=1-p``.
|
|
|
|
The probability mass function of this distribution, over possible outcomes ``k``, is
|
|
|
|
.. math::
|
|
|
|
{\begin{cases}
|
|
q=1-p & \text{if }value=0 \\
|
|
p & \text{if }value=1
|
|
\end{cases}}
|
|
|
|
Args:
|
|
probs (float|Tensor): The ``probs`` input of Bernoulli distribution. The data type is float32 or float64. The range must be in [0, 1].
|
|
name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> # init `probs` with a float
|
|
>>> rv = Bernoulli(probs=0.3)
|
|
|
|
>>> print(rv.mean)
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
0.30000001)
|
|
|
|
>>> print(rv.variance)
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
0.21000001)
|
|
|
|
>>> print(rv.entropy())
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
0.61086434)
|
|
"""
|
|
|
|
name: str
|
|
probs: Tensor
|
|
logits: Tensor
|
|
dtype: _DTypeLiteral
|
|
|
|
def __init__(self, probs: float | Tensor, name: str | None = None) -> None:
|
|
self.name = name or 'Bernoulli'
|
|
if not in_dynamic_mode():
|
|
check_type(
|
|
probs,
|
|
'probs',
|
|
(float, Variable, paddle.pir.Value),
|
|
self.name,
|
|
)
|
|
|
|
# Get/convert probs to tensor.
|
|
if self._validate_args(probs):
|
|
self.probs = probs
|
|
self.dtype = convert_dtype(probs.dtype)
|
|
else:
|
|
[self.probs] = self._to_tensor(probs)
|
|
self.dtype = paddle.get_default_dtype()
|
|
|
|
# Clip probs from [0, 1] to (0, 1) with smallest representable number `eps`.
|
|
self.probs = _clip_probs(self.probs, self.dtype)
|
|
self.logits = self._probs_to_logits(self.probs, is_binary=True)
|
|
|
|
super().__init__(batch_shape=self.probs.shape, event_shape=())
|
|
|
|
@property
|
|
def mean(self) -> Tensor:
|
|
"""Mean of Bernoulli distribution.
|
|
|
|
Returns:
|
|
Tensor: Mean value of distribution.
|
|
"""
|
|
return self.probs
|
|
|
|
@property
|
|
def variance(self) -> Tensor:
|
|
"""Variance of Bernoulli distribution.
|
|
|
|
Returns:
|
|
Tensor: Variance value of distribution.
|
|
"""
|
|
return paddle.multiply(self.probs, (1 - self.probs))
|
|
|
|
@param_one_alias(["shape", "sample_shape"])
|
|
def sample(self, shape: Sequence[int] = []) -> Tensor:
|
|
"""Sample from Bernoulli distribution.
|
|
|
|
Args:
|
|
shape (Sequence[int], optional): Sample shape.
|
|
|
|
Returns:
|
|
Tensor: Sampled data with shape `sample_shape` + `batch_shape` + `event_shape`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(paddle.full([1], 0.3))
|
|
>>> print(rv.sample([100]).shape)
|
|
paddle.Size([100, 1])
|
|
|
|
>>> rv = Bernoulli(paddle.to_tensor(0.3))
|
|
>>> print(rv.sample([100]).shape)
|
|
paddle.Size([100])
|
|
|
|
>>> rv = Bernoulli(paddle.to_tensor([0.3, 0.5]))
|
|
>>> print(rv.sample([100]).shape)
|
|
paddle.Size([100, 2])
|
|
|
|
>>> rv = Bernoulli(paddle.to_tensor([0.3, 0.5]))
|
|
>>> print(rv.sample([100, 2]).shape)
|
|
paddle.Size([100, 2, 2])
|
|
"""
|
|
name = self.name + '_sample'
|
|
if not in_dynamic_mode():
|
|
check_type(
|
|
shape,
|
|
'shape',
|
|
(np.ndarray, Variable, list, tuple, paddle.pir.Value),
|
|
name,
|
|
)
|
|
|
|
shape = shape if isinstance(shape, tuple) else tuple(shape)
|
|
shape = self._extend_shape(shape)
|
|
|
|
with paddle.no_grad():
|
|
return paddle.bernoulli(self.probs.expand(shape), name=name)
|
|
|
|
@param_one_alias(["shape", "sample_shape"])
|
|
def rsample(
|
|
self, shape: Sequence[int] = [], temperature: float = 1.0
|
|
) -> Tensor:
|
|
"""Sample from Bernoulli distribution (reparameterized).
|
|
|
|
The `rsample` is a continuously approximate of Bernoulli distribution reparameterized sample method.
|
|
[1] Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables. 2016.
|
|
[2] Eric Jang, Shixiang Gu, and Ben Poole. Categorical Reparameterization with Gumbel-Softmax. 2016.
|
|
|
|
Note:
|
|
`rsample` need to be followed by a `sigmoid`, which converts samples' value to unit interval (0, 1).
|
|
|
|
Args:
|
|
shape (Sequence[int], optional): Sample shape.
|
|
temperature (float): temperature for rsample, must be positive.
|
|
|
|
Returns:
|
|
Tensor: Sampled data with shape `sample_shape` + `batch_shape` + `event_shape`.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> paddle.seed(1)
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(paddle.full([1], 0.3))
|
|
>>> print(rv.sample([100]).shape)
|
|
paddle.Size([100, 1])
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> print(rv.rsample([100]).shape)
|
|
paddle.Size([100])
|
|
|
|
>>> rv = Bernoulli(paddle.to_tensor([0.3, 0.5]))
|
|
>>> print(rv.rsample([100]).shape)
|
|
paddle.Size([100, 2])
|
|
|
|
>>> rv = Bernoulli(paddle.to_tensor([0.3, 0.5]))
|
|
>>> print(rv.rsample([100, 2]).shape)
|
|
paddle.Size([100, 2, 2])
|
|
|
|
>>> # `rsample` has to be followed by a `sigmoid`
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> rsample = rv.rsample([3])
|
|
>>> rsample_sigmoid = paddle.nn.functional.sigmoid(rsample)
|
|
>>> print(rsample)
|
|
Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[-1.46112013, -0.01239836, -1.32765460])
|
|
>>> print(rsample_sigmoid)
|
|
Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[0.18829606, 0.49690047, 0.20954758])
|
|
|
|
>>> # The smaller the `temperature`, the distribution of `rsample` closer to `sample`, with `probs` of 0.3.
|
|
>>> print(
|
|
... paddle.nn.functional.sigmoid(
|
|
... rv.rsample(
|
|
... [1000],
|
|
... temperature=1.0,
|
|
... )
|
|
... ).sum()
|
|
... )
|
|
>>> # doctest: +SKIP('output will be different')
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
365.63122559)
|
|
>>> # doctest: -SKIP
|
|
|
|
>>> print(
|
|
... paddle.nn.functional.sigmoid(
|
|
... rv.rsample(
|
|
... [1000],
|
|
... temperature=0.1,
|
|
... )
|
|
... ).sum()
|
|
... )
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
320.15057373)
|
|
"""
|
|
name = self.name + '_rsample'
|
|
if not in_dynamic_mode():
|
|
check_type(
|
|
shape,
|
|
'shape',
|
|
(np.ndarray, Variable, paddle.pir.Value, list, tuple),
|
|
name,
|
|
)
|
|
check_type(
|
|
temperature,
|
|
'temperature',
|
|
(float,),
|
|
name,
|
|
)
|
|
|
|
shape = shape if isinstance(shape, tuple) else tuple(shape)
|
|
shape = self._extend_shape(shape)
|
|
|
|
temperature = paddle.full(
|
|
shape=(), fill_value=temperature, dtype=self.dtype
|
|
)
|
|
|
|
probs = self.probs.expand(shape)
|
|
uniforms = paddle.rand(shape, dtype=self.dtype)
|
|
return paddle.divide(
|
|
paddle.add(
|
|
paddle.subtract(uniforms.log(), (-uniforms).log1p()),
|
|
paddle.subtract(probs.log(), (-probs).log1p()),
|
|
),
|
|
temperature,
|
|
)
|
|
|
|
def cdf(self, value: Tensor) -> Tensor:
|
|
r"""Cumulative distribution function(CDF) evaluated at value.
|
|
|
|
.. math::
|
|
|
|
{ \begin{cases}
|
|
0 & \text{if } value \lt 0 \\
|
|
1 - p & \text{if } 0 \leq value \lt 1 \\
|
|
1 & \text{if } value \geq 1
|
|
\end{cases}
|
|
}
|
|
|
|
Args:
|
|
value (Tensor): Value to be evaluated.
|
|
|
|
Returns:
|
|
Tensor: CDF evaluated at value.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> print(rv.cdf(paddle.to_tensor([1.0])))
|
|
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[1.])
|
|
"""
|
|
name = self.name + '_cdf'
|
|
if not in_dynamic_mode():
|
|
check_type(value, 'value', (Variable, paddle.pir.Value), name)
|
|
|
|
value = self._check_values_dtype_in_probs(self.probs, value)
|
|
probs, value = paddle.broadcast_tensors([self.probs, value])
|
|
|
|
zeros = paddle.zeros_like(probs)
|
|
ones = paddle.ones_like(probs)
|
|
|
|
return paddle.where(
|
|
value < 0,
|
|
zeros,
|
|
paddle.where(value < 1, paddle.subtract(ones, probs), ones),
|
|
name=name,
|
|
)
|
|
|
|
def log_prob(self, value: Tensor) -> Tensor:
|
|
"""Log of probability density function.
|
|
|
|
Args:
|
|
value (Tensor): Value to be evaluated.
|
|
|
|
Returns:
|
|
Tensor: Log of probability density evaluated at value.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> print(rv.log_prob(paddle.to_tensor([1.0])))
|
|
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[-1.20397282])
|
|
"""
|
|
name = self.name + '_log_prob'
|
|
if not in_dynamic_mode():
|
|
check_type(value, 'value', (Variable, paddle.pir.Value), name)
|
|
|
|
value = self._check_values_dtype_in_probs(self.probs, value)
|
|
logits, value = paddle.broadcast_tensors([self.logits, value])
|
|
return -binary_cross_entropy_with_logits(
|
|
logits, value, reduction='none', name=name
|
|
)
|
|
|
|
def prob(self, value: Tensor) -> Tensor:
|
|
r"""Probability density function(PDF) evaluated at value.
|
|
|
|
.. math::
|
|
|
|
{ \begin{cases}
|
|
q=1-p & \text{if }value=0 \\
|
|
p & \text{if }value=1
|
|
\end{cases}
|
|
}
|
|
|
|
Args:
|
|
value (Tensor): Value to be evaluated.
|
|
|
|
Returns:
|
|
Tensor: PDF evaluated at value.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> print(rv.prob(paddle.to_tensor([1.0])))
|
|
Tensor(shape=[1], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
[0.29999998])
|
|
"""
|
|
name = self.name + '_prob'
|
|
if not in_dynamic_mode():
|
|
check_type(value, 'value', (Variable, paddle.pir.Value), name)
|
|
|
|
return self.log_prob(value).exp(name=name)
|
|
|
|
def entropy(self) -> Tensor:
|
|
r"""Entropy of Bernoulli distribution.
|
|
|
|
.. math::
|
|
|
|
{
|
|
entropy = -(q \log q + p \log p)
|
|
}
|
|
|
|
Returns:
|
|
Tensor: Entropy of distribution.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> print(rv.entropy())
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
0.61086434)
|
|
"""
|
|
name = self.name + '_entropy'
|
|
|
|
return binary_cross_entropy_with_logits(
|
|
self.logits, self.probs, reduction='none', name=name
|
|
)
|
|
|
|
def kl_divergence(self, other: Bernoulli) -> Tensor:
|
|
r"""The KL-divergence between two Bernoulli distributions.
|
|
|
|
.. math::
|
|
|
|
{
|
|
KL(a || b) = p_a \log(p_a / p_b) + (1 - p_a) \log((1 - p_a) / (1 - p_b))
|
|
}
|
|
|
|
Args:
|
|
other (Bernoulli): instance of Bernoulli.
|
|
|
|
Returns:
|
|
Tensor: kl-divergence between two Bernoulli distributions.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> from paddle.distribution import Bernoulli
|
|
|
|
>>> rv = Bernoulli(0.3)
|
|
>>> rv_other = Bernoulli(0.7)
|
|
|
|
>>> print(rv.kl_divergence(rv_other))
|
|
Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
|
|
0.33891910)
|
|
"""
|
|
name = self.name + '_kl_divergence'
|
|
if not in_dynamic_mode():
|
|
check_type(other, 'other', Bernoulli, name)
|
|
|
|
a_logits = self.logits
|
|
b_logits = other.logits
|
|
|
|
log_pa = -softplus(-a_logits)
|
|
log_pb = -softplus(-b_logits)
|
|
|
|
pa = sigmoid(a_logits)
|
|
one_minus_pa = sigmoid(-a_logits)
|
|
|
|
log_one_minus_pa = -softplus(a_logits)
|
|
log_one_minus_pb = -softplus(b_logits)
|
|
|
|
return paddle.add(
|
|
paddle.subtract(
|
|
paddle.multiply(log_pa, pa), paddle.multiply(log_pb, pa)
|
|
),
|
|
paddle.subtract(
|
|
paddle.multiply(log_one_minus_pa, one_minus_pa),
|
|
paddle.multiply(log_one_minus_pb, one_minus_pa),
|
|
),
|
|
)
|