mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
80fa8e1007
Summary: This adds a `.expand` method for distributions that is akin to the `torch.Tensor.expand` method for tensors. It returns a new distribution instance with batch dimensions expanded to the desired `batch_shape`. Since this calls `torch.Tensor.expand` on the distribution's parameters, it does not allocate new memory for the expanded distribution instance's parameters. e.g. ```python >>> d = dist.Normal(torch.zeros(100, 1), torch.ones(100, 1)) >>> d.sample().shape torch.Size([100, 1]) >>> d.expand([100, 10]).sample().shape torch.Size([100, 10]) ``` We have already been using the `.expand` method in Pyro in our [patch](https://github.com/uber/pyro/blob/dev/pyro/distributions/torch.py#L10) of `torch.distributions`. We use this in our models to enable dynamic broadcasting. This has also been requested by a few users on the distributions slack, and we believe will be useful to the larger community. Note that currently, there is no convenient and efficient way to expand distribution instances: - Many distributions use `TransformedDistribution` (or wrap over another distribution instance. e.g. `OneHotCategorical` uses a `Categorical` instance) under the hood, or have lazy parameters. This makes it difficult to collect all the relevant parameters, broadcast them and construct new instances. - In the few cases where this is even possible, the resulting implementation would be inefficient since we will go through a lot of broadcasting and args validation logic in `__init__.py` that can be avoided. The `.expand` method allows for a safe and efficient way to expand distribution instances. Additionally, this bypasses `__init__.py` (using `__new__` and populating relevant attributes) since we do not need to do any broadcasting or args validation (which was already done when the instance was first created). This can result in significant savings as compared to constructing new instances via `__init__` (that said, the `sample` and `log_prob` methods will probably be the rate determining steps in many applications). e.g. ```python >>> a = dist.Bernoulli(torch.ones([10000, 1]), validate_args=True) >>> %timeit a.expand([10000, 100]) 15.2 µs ± 224 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) >>> %timeit dist.Bernoulli(torch.ones([10000, 100]), validate_args=True) 11.8 ms ± 153 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) ``` cc. fritzo, apaszke, vishwakftw, alicanb Pull Request resolved: https://github.com/pytorch/pytorch/pull/11341 Differential Revision: D9728485 Pulled By: soumith fbshipit-source-id: 3b94c23bc6a43ee704389e6287aa83d1e278d52f
97 lines
3.3 KiB
Python
97 lines
3.3 KiB
Python
import math
|
|
from numbers import Number
|
|
|
|
import torch
|
|
from torch.distributions import constraints
|
|
from torch.distributions.exp_family import ExponentialFamily
|
|
from torch.distributions.utils import broadcast_all
|
|
|
|
|
|
class Normal(ExponentialFamily):
|
|
r"""
|
|
Creates a normal (also called Gaussian) distribution parameterized by
|
|
:attr:`loc` and :attr:`scale`.
|
|
|
|
Example::
|
|
|
|
>>> m = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
|
|
>>> m.sample() # normally distributed with loc=0 and scale=1
|
|
tensor([ 0.1046])
|
|
|
|
Args:
|
|
loc (float or Tensor): mean of the distribution (often referred to as mu)
|
|
scale (float or Tensor): standard deviation of the distribution
|
|
(often referred to as sigma)
|
|
"""
|
|
arg_constraints = {'loc': constraints.real, 'scale': constraints.positive}
|
|
support = constraints.real
|
|
has_rsample = True
|
|
_mean_carrier_measure = 0
|
|
|
|
@property
|
|
def mean(self):
|
|
return self.loc
|
|
|
|
@property
|
|
def stddev(self):
|
|
return self.scale
|
|
|
|
@property
|
|
def variance(self):
|
|
return self.stddev.pow(2)
|
|
|
|
def __init__(self, loc, scale, validate_args=None):
|
|
self.loc, self.scale = broadcast_all(loc, scale)
|
|
if isinstance(loc, Number) and isinstance(scale, Number):
|
|
batch_shape = torch.Size()
|
|
else:
|
|
batch_shape = self.loc.size()
|
|
super(Normal, self).__init__(batch_shape, validate_args=validate_args)
|
|
|
|
def expand(self, batch_shape, _instance=None):
|
|
new = self._get_checked_instance(Normal, _instance)
|
|
batch_shape = torch.Size(batch_shape)
|
|
new.loc = self.loc.expand(batch_shape)
|
|
new.scale = self.scale.expand(batch_shape)
|
|
super(Normal, new).__init__(batch_shape, validate_args=False)
|
|
new._validate_args = self._validate_args
|
|
return new
|
|
|
|
def sample(self, sample_shape=torch.Size()):
|
|
shape = self._extended_shape(sample_shape)
|
|
with torch.no_grad():
|
|
return torch.normal(self.loc.expand(shape), self.scale.expand(shape))
|
|
|
|
def rsample(self, sample_shape=torch.Size()):
|
|
shape = self._extended_shape(sample_shape)
|
|
eps = self.loc.new(shape).normal_()
|
|
return self.loc + eps * self.scale
|
|
|
|
def log_prob(self, value):
|
|
if self._validate_args:
|
|
self._validate_sample(value)
|
|
# compute the variance
|
|
var = (self.scale ** 2)
|
|
log_scale = math.log(self.scale) if isinstance(self.scale, Number) else self.scale.log()
|
|
return -((value - self.loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
|
|
|
|
def cdf(self, value):
|
|
if self._validate_args:
|
|
self._validate_sample(value)
|
|
return 0.5 * (1 + torch.erf((value - self.loc) * self.scale.reciprocal() / math.sqrt(2)))
|
|
|
|
def icdf(self, value):
|
|
if self._validate_args:
|
|
self._validate_sample(value)
|
|
return self.loc + self.scale * torch.erfinv(2 * value - 1) * math.sqrt(2)
|
|
|
|
def entropy(self):
|
|
return 0.5 + 0.5 * math.log(2 * math.pi) + torch.log(self.scale)
|
|
|
|
@property
|
|
def _natural_params(self):
|
|
return (self.loc / self.scale.pow(2), -0.5 * self.scale.pow(2).reciprocal())
|
|
|
|
def _log_normalizer(self, x, y):
|
|
return -0.25 * x.pow(2) / y + 0.5 * torch.log(-math.pi / y)
|