PyTorch torch.distributions

A distribution is an object that describes how random values are spread. torch.distributions builds distributions that can sample values and score them.

A Normal distribution

Normal is the bell-shaped distribution, set by a mean and a standard deviation. The standard deviation, stored as scale, measures the spread around the mean.

torch.manual_seed(0)
dist = Normal(2.0, 0.5)
print("Mean:", dist.mean.item())
print("Scale:", dist.scale.item())
print("One sample:", round(dist.sample().item(), 4))

Score a value, measure the spread

log_prob scores how typical one value is under the distribution. Entropy is one number for how spread out the distribution is.

value = torch.tensor(2.5)
print("log_prob of 2.5:", round(dist.log_prob(value).item(), 4))
print("Entropy:", round(dist.entropy().item(), 4))

A Categorical distribution

Categorical picks one class from a list of probabilities. The probabilities in the list add up to 1.

torch.manual_seed(1)
cat = Categorical(probs=torch.tensor([0.5, 0.3, 0.2]))
print("Probs:", [round(v, 2) for v in cat.probs.tolist()])
print("One pick:", cat.sample().item())
print("Entropy:", round(cat.entropy().item(), 4))

Example

import torch
from torch.distributions import Normal

torch.manual_seed(0)
dist = Normal(2.0, 0.5)
samples = dist.sample((4,))
log_probs = dist.log_prob(samples)

for value, lp in zip(samples.tolist(), log_probs.tolist()):
    print(f"Sample {value:.4f}  log_prob {lp:.4f}")
print("Entropy:", round(dist.entropy().item(), 4))

Key points

Distribution: an object that describes how random values spread
Use Normal(mean, std) to build a bell-shaped distribution around a mean
Use dist.sample(shape) to draw random values from a distribution
Use dist.log_prob(value) to score how typical a value is
Use dist.entropy() to get one spread number for a distribution
One distribution object draws values, scores values, and reports its spread

In this example, a Normal distribution is built with mean 2.0 and scale 0.5. The seed is set before sampling, so the four draws repeat on every run. Each draw prints with its log_prob score beside it. The sample 0.9106 sits far from the mean, and its score is the lowest at -2.5994. The sample 1.8533 sits near the mean, and its score is the highest at -0.2688. The last line prints the entropy of the distribution, 0.7258.

Note

Draws from sample are random, so a seed decides which values appear

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch.distributions import Normal

torch.manual_seed(0)
dist = Normal(2.0, 0.5)
samples = dist.sample((4,))
log_probs = dist.log_prob(samples)

for value, lp in zip(samples.tolist(), log_probs.tolist()):
    print(f"Sample {value:.4f}  log_prob {lp:.4f}")
print("Entropy:", round(dist.entropy().item(), 4))

Exercise

Which method scores how typical a value is under a distribution?

sample log_prob entropy argmax