Pyro 中分布具备维度概念,并且随机函数的维度不是数据的维度!并且非常重要的事情是你需要把所有的数据看成某“一个”随机向量的样本。
官方中文文档
causalai.github.io/pyro_zh_tutorial/tensor_shapes.html
Pyro中随机函数的维度
本教程介绍 Pyro‘s organization of tensor dimensions. 开始之前,您应该熟悉PyTorch 广播语义, see PyTorch broadcasting semantics.
并行枚举是本节最难懂的内容,它到底如何理解?
在 Pyro 中,正确地理解和使用维度概念是至关重要的,因为它涉及到如何构建和解释概率模型。以下是对您提供的关键点的解释和理解:
随机函数的维度:与 PyTorch 中的 Tensor 不同,Pyro 中的 Distribution 有两个维度属性:
batch_shape
和event_shape
。这两个属性定义了样本的维度。
batch_shape
表示不同参数的同一类型随机变量数量的维度,意味着随机变量之间的条件独立性。event_shape
是随机变量的一个事件所具有的维度,意味着随机向量分量之间的相依性。样本的维度:一个样本的维度由
batch_shape
和event_shape
相加得出。当你从分布中采样时,样本的形状将包括这两个维度。广播(Broadcasting):在 Pyro 中,张量操作遵循 PyTorch 的广播规则,这意味着在进行操作时,形状不匹配的张量可以根据规则自动扩展以匹配形状。
使用
.expand()
和plate
:.expand()
方法可以用来创建一批相同分布的样本,而plate
上下文管理器用来声明一组条件独立的变量,Pyro 会自动扩展这些变量以匹配指定的批次大小。
.to_event(n)
方法:这个方法允许你将分布的右起n
个维度声明为 dependent,从而改变随机函数的形状。这对于处理多维数据时特别有用。
pyro.plate
:使用with pyro.plate('name', size):
可以声明一个条件独立的维度,Pyro 会利用这一点来进行有效的并行计算。维度声明:在 Pyro 中,所有的维度必须被明确地声明为 dependent(依赖的)或者是 conditionally independent(条件独立的)。
支持左侧批量处理:尽量在左侧支持批量处理,这可以让 Pyro 自动并行化计算。
使用负索引和省略号:在访问张量时,使用负索引(如
x.sum(-1)
)和省略号(如pixel = image[..., i, j]
)可以更灵活地处理多维数据。调试:当调试模型时,可以使用
Trace.format_shapes()
方法来检查追踪中每个样本点的形状,这有助于确保维度的正确性。理解这些概念对于在 Pyro 中构建和调试概率模型非常重要。正确地声明和使用维度可以帮助你更有效地构建模型,并确保计算的正确性。如果有任何具体的模型或代码段需要帮助,请随时提问。
学完本文,您将学会:
-
训练模型和 debug 的时候,设置
pyro.enable_validation(True)
. -
张量按照最右边的维度进行 Broadcasting:
torch.ones(3,4,5) + torch.ones(5)
. -
样本维度 =
batch_shape + event_shape
. -
对于样本 x�
model.log_prob(x).shape == batch_shape
-
用
.expand()
抽取一批量的样本,或者用plate
自动 expand. -
使用
my_dist.to_event(n)
将右起 n 个维度声明为 dependent,从而是实现随机函数 reshape. -
使用
with pyro.plate('name', size):
来申明某个维度条件独立性。 -
所有的维度必须被申明为 dependent 或者条件独立.
-
Try to support batching on the left. This lets Pyro auto-parallelize.
-
use negative indices like
x.sum(-1)
rather thanx.sum(2)
-
use ellipsis notation like
pixel = image[..., i, j]
-
use Vindex if
i,j
are enumerated,pixel = Vindex(image)[..., i, j]
-
-
When debugging, examine all shapes in a trace using Trace.format_shapes().
[1]:
import os, torch, pyro from torch.distributions import constraints from pyro.distributions import Bernoulli, Categorical, MultivariateNormal, Normal from pyro.distributions.util import broadcast_shape from pyro.infer import Trace_ELBO, TraceEnum_ELBO, config_enumerate import pyro.poutine as poutine from pyro.optim import Adam smoke_test = ('CI' in os.environ) assert pyro.__version__.startswith('1.3.0') pyro.enable_validation(True) # <---- This is always a good idea! # We'll ue this helper to check our models are correct. def test_model(model, guide, loss): pyro.clear_param_store() loss.loss(model, guide)
随机函数的维度
在 Pyro 这个概率编程库中,对维度的处理与 PyTorch 中 Tensor 的处理有所不同。以下是对 Pyro 中随机函数维度概念的解释:
1. **Tensor vs. Distribution 的维度**:
- 在 PyTorch 中,Tensor 只有一个表示其形状的维度属性 `.shape`。
- 在 Pyro 中,Distribution 有两个维度属性:`.batch_shape` 和 `.event_shape`。这两个属性共同定义了样本的维度。2. **样本的维度**:
- 当你从分布 `d` 中采样得到 `x` 时,`x` 的形状应该是分布的 `batch_shape` 和 `event_shape` 的总和,即 `x.shape == d.batch_shape + d.event_shape`。3. **数学上的解释**:
- 随机变量 \( X \sim p(x; \theta) \) 的某个事件可以表示为 \( X = x \),这里的维度就是 `event_shape`。
- 如果我们需要一批具有不同参数 \( \theta \) 的随机变量 \( X \)(称为 batch r.v.),那么 `batch_shape` 就是这些不同参数的随机变量数量的维度。4. **条件独立性与相依性**:
- 在 `.batch_shape` 上的索引意味着随机变量之间的条件独立性。
- 在 `.event_shape` 上的索引意味着随机向量分量之间的相依性,即从同一分布中抽取的样本之间的关系。5. **`.log_prob()` 方法**:
- 由于一个随机向量样本对应一个概率值,`.log_prob()` 方法为每个形状为 `.event_shape` 的事件生成一个概率值。因此,`.log_prob()` 方法的输出维度就是 `.batch_shape`。6. **`.sample()` 方法的 `sample_shape` 参数**:
- `Distribution.sample()` 方法还有一个参数 `sample_shape`,它表示从分布中抽取独立同分布样本的维度。7. **示例**:
- 假设一个班级中每个学生的身高和体重组成一个随机向量,那么 `event_shape = (2,)`。
- 如果一个学校有多个这样的班级,每个班级有其独特的身高和体重分布参数,那么 `batch_shape` 就是班级的数量。
- 如果从某个总体中抽取样本,那么 `sample_shape` 就是从该总体中抽取的样本数量。通过这种方式,Pyro 允许你以一种非常灵活和表达性强的方式来构建和操作概率模型,同时保持对维度的精确控制。这对于处理复杂的贝叶斯模型和进行概率推断至关重要。
Pytorch 中的 Tensor
只有一个维度属性 .shape
,但是 Pyro 中 Distribution
有两个维度属性: .batch_shape
和 .event_shape
. 他们两个一起定义了一个样本的维度。
x = d.sample() assert x.shape == d.batch_shape + d.event_shape # 两个元组的并,不要误解成对应每个维度相加
数学上来说,随机变量 X∼p(x;θ)�∼�(�;�) 的某个事件是 X=x�=�,因此维度就是 event_shape
。而我们需要一个 batch 不同参数 θ� 的 r.v. X� (称之为 batch r.v.),所以 batch_shape
就是不同参数同一类型的随机变量数量的维度。于是最后的一个 “batch r.v.(batch of 同种类型随机变量)” 样本的维度就是 batch_shape
并上 event_shape
.
Indices over .batch_shape
意味着随机变量之间的条件独立性, 而 indices over .event_shape
意味着随机向量分量之间的相依性 (即从分布中抽出一个样本). 因为一个随机向量样本对应这一个概率值, 所以 .log_prob()
方法为 each event of shape .event_shape
产生一个概率值,所以 .log_prob()
方法输出的维度就是 .batch_shape
:
# 每个样本对应一个概率值,batch_shape 就是 batch r.v. 的维度。 assert d.log_prob(x).shape == d.batch_shape
请注意 Distribution.sample()
方法也有参数 sample_shape
, 它用于表示batch r.v. 的独立同分布采样维度,所以
x2 = d.sample(sample_shape) assert x2.shape == sample_shape + batch_shape + event_shape
也就是说 Pyro 中每个分布(也就是随机函数)都具备两个维度 batch_shape 和 event_shape,抽样之后我们还可以得到一个 sample_shape
。
我们来举个例子,某个班级每个学生身高和体重组成一个随机向量,那么 event_shape = (2, )
;而一个学校有多个不同的班级,其中每个班级具备不同的身高体重分布参数,那么 batch_shape = 班级数量
;最后这个学校全体学生来自于某个总体,那么此时从该总体中可以抽样得到 sample_shape
个样本。
In summary
| iid | independent | dependent ------+--------------+-------------+------------ shape = sample_shape + batch_shape + event_shape
简单例子
最简单的分布形状是一个单变量分布。
[5]:
d = Bernoulli(0.5) assert d.batch_shape == () assert d.event_shape == () x = d.sample() assert x.shape == () assert d.log_prob(x).shape == ()
通过传入批量参数可得到 batched 分布(batch r.v.)。
[6]:
d = Bernoulli(0.5 * torch.ones(3,4)) assert d.batch_shape == (3, 4) assert d.event_shape == () x = d.sample() assert x.shape == (3, 4) assert d.log_prob(x).shape == (3, 4)
另外一种 batch distributions 的方式是使用 .expand()
method(该方法要求最右边的维度要相等).
[2]:
d = Bernoulli(torch.tensor([0.1, 0.2, 0.3, 0.4])).expand([3, 4]) assert d.batch_shape == (3, 4) assert d.event_shape == () x = d.sample() assert x.shape == (3, 4) assert d.log_prob(x).shape == (3, 4)
多元分布具有非空 .event_shape
. 对于这些分布, .sample()
和 .log_prob(x)
的维度会不一样:
[8]:
d = MultivariateNormal(torch.zeros(3), torch.eye(3, 3)) assert d.batch_shape == () assert d.event_shape == (3,) x = d.sample() assert x.shape == (3,) # == batch_shape + event_shape assert d.log_prob(x).shape == () # == batch_shape
下一小节提示:
x = pyro.sample("x", dist.Normal(0, 1).expand([10]).to_event(1)) with pyro.plate("x_plate", 10): x = pyro.sample("x", dist.Normal(0, 1))
变换event_shape
在 Pyro 中,您可以使用 .to_event(n) 将一维分布转化成多维分布(也就是随便变量 stack 成随机向量),其中 n
表示从右边开始数 n 个维度,声明 dependent
(也就是他们为随机向量). 所以 to_event(0)
相当于分布没有维度变换,而 to_event(1)
表示随机向量的维度为原来分布 batch_shape
的最后一个维度。
[24]:
d = Bernoulli(0.5 * torch.ones(3,4)).to_event(1) assert d.batch_shape == (3,) assert d.event_shape == (4,) x = d.sample() assert x.shape == (3, 4) assert d.log_prob(x).shape == (3,)
在使用Pyro程序时,请记住样本具有维度 batch_shape + event_shape
, 而 .log_prob(x)
具有维度batch_shape
. 你需要确保 batch_shape
is carefully controlled by either trimming it down with .to_event(n)
or by declaring dimensions as independent via pyro.plate
.
实际应用中:It is always safe to assume dependence。在 Pyro 中我们经常将某些维度声明为 dependent,即使他们实际上是独立的, e.g.
x = pyro.sample("x", dist.Normal(0, 1).expand([10]).to_event(1)) assert x.shape == (10,)
这很有用,原因有两个:
-
(容易构造多元分布) it allows us to easily swap in a
MultivariateNormal
distribution later. -
(简化代码,避免使用
plate
) it simplifies the code a bit since we don’t need aplate
(see below) as in
with pyro.plate("x_plate", 10): x = pyro.sample("x", dist.Normal(0, 1)) # .expand([10]) is automatic assert x.shape == (10,)
这两个版本之间的区别是: 第二个有 plate
的版本告知 Pyro 使用条件独立性信息 when estimating gradients, 而第一个版本 Pyro 必须假定分量之间是相关的 (even though the normals are in fact conditionally independent). 这类似于图模型中的 d�-分离: it is always safe to add edges and assume variables may be dependent (i.e. to widen the model class), but it is unsafe to assume independence when variables are actually dependent (i.e. narrowing the model class so the true model lies outside of the class, as in mean field).
实际上,Pyro的 SVI 推理算法对 Normal
distributions 使用了重参数化技巧,因此两个梯度估计量具有相同的性能。
下节内容提示:
x_axis = pyro.plate("x_axis", 3, dim=-2) y_axis = pyro.plate("y_axis", 2, dim=-3) with x_axis: x = pyro.sample("x", Normal(0, 1)) with y_axis: y = pyro.sample("y", Normal(0, 1)) with x_axis, y_axis: xy = pyro.sample("xy", Normal(0, 1)) z = pyro.sample("z", Normal(0, 1).expand([5]).to_event(1)) assert x.shape == (3, 1) # batch_shape == (3,1) event_shape == () assert y.shape == (2, 1, 1) # batch_shape == (2,1,1) event_shape == () assert xy.shape == (2, 3, 1) # batch_shape == (2,3,1) event_shape == () assert z.shape == (2, 3, 1, 5) # batch_shape == (2,3,1) event_shape == (5,)
变换batch_shape
用 plate
声明条件独立性来变换 batch_shape
.
Pyro模型可以使用上下文管理器 pyro.plate 来申明特定批量维度(batch dimensions) 是独立的。然后,推理算法可以利用这种独立性,e.g. construct lower variance gradient estimators or to enumerate in linear space rather than exponential space. An example of an independent dimension is the index over data in a minibatch: each datum should be independent of all others.
最简单的方式申明独立行的方法如下:
with pyro.plate("my_plate"): # within this context, batch dimension -1 is independent
我们推荐提供一个 size
参数来帮助 debug shapes:
with pyro.plate("my_plate", len(my_data)): # within this context, batch dimension -1 is independent
从 Pyro0.2开始,您还可以嵌套 plate
, e.g. if you have per-pixel independence:
with pyro.plate("x_axis", 320): # within this context, batch dimension -1 is independent with pyro.plate("y_axis", 200): # within this context, batch dimensions -2 and -1 are independent
注意我们总是用负数 -2, -1 来表示从右边索引。
Finally if you want to mix and match plate
s for e.g. noise that depends only on x
, some noise that depends only on y
, and some noise that depends on both, you can declare multiple plates
and use them as reusable context managers. In this case Pyro cannot automatically allocate a dimension, so you need to provide a dim
argument (again counting from the right):
x_axis = pyro.plate("x_axis", 3, dim=-2) y_axis = pyro.plate("y_axis", 2, dim=-3) with x_axis: # within this context, batch dimension -2 is independent with y_axis: # within this context, batch dimension -3 is independent with x_axis, y_axis: # within this context, batch dimensions -3 and -2 are independent
Let’s take a closer look at batch sizes within plate
s.
[73]:
def model1(): a = pyro.sample("a", Normal(0, 1)) b = pyro.sample("b", Normal(torch.zeros(2), 1).to_event(1)) with pyro.plate("c_plate", 2): c = pyro.sample("c", Normal(torch.zeros(2), 1)) with pyro.plate("d_plate", 3): d = pyro.sample("d", Normal(torch.zeros(3,4,5), 1).to_event(2)) assert a.shape == () # batch_shape == () event_shape == () assert b.shape == (2,) # batch_shape == () event_shape == (2,) assert c.shape == (2,) # batch_shape == (2,) event_shape == () assert d.shape == (3,4,5) # batch_shape == (3,) event_shape == (4,5) x_axis = pyro.plate("x_axis", 3, dim=-2) y_axis = pyro.plate("y_axis", 2, dim=-3) with x_axis: x = pyro.sample("x", Normal(0, 1)) with y_axis: y = pyro.sample("y", Normal(0, 1)) with x_axis, y_axis: xy = pyro.sample("xy", Normal(0, 1)) z = pyro.sample("z", Normal(0, 1).expand([5]).to_event(1)) assert x.shape == (3, 1) # batch_shape == (3,1) event_shape == () assert y.shape == (2, 1, 1) # batch_shape == (2,1,1) event_shape == () assert xy.shape == (2, 3, 1) # batch_shape == (2,3,1) event_shape == () assert z.shape == (2, 3, 1, 5) # batch_shape == (2,3,1) event_shape == (5,) test_model(model1, model1, Trace_ELBO())
It is helpful to visualize the .shape
s of each sample site by aligning them at the boundary between batch_shape
and event_shape
: dimensions to the right will be summed out in .log_prob()
and dimensions to the left will remain.
batch dims | event dims -----------+----------- | a = sample("a", Normal(0, 1)) |2 b = sample("b", Normal(zeros(2), 1) | .to_event(1)) | with plate("c", 2): 2| c = sample("c", Normal(zeros(2), 1)) | with plate("d", 3): 3|4 5 d = sample("d", Normal(zeros(3,4,5), 1) | .to_event(2)) | | x_axis = plate("x", 3, dim=-2) | y_axis = plate("y", 2, dim=-3) | with x_axis: 3 1| x = sample("x", Normal(0, 1)) | with y_axis: 2 1 1| y = sample("y", Normal(0, 1)) | with x_axis, y_axis: 2 3 1| xy = sample("xy", Normal(0, 1)) 2 3 1|5 z = sample("z", Normal(0, 1).expand([5]) | .to_event(1))
为了自动的检查程序中样本的维度,you can trace the program 并且使用方法 Trace.format_shapes() 来打印每个 sample site 的三种维度:
-
分布的维度(both
site["fn"].batch_shape
andsite["fn"].event_shape
), -
the value shape(
site["value"].shape
), -
如果对数概率被计算来,那么也返回
log_prob
维度 (site["log_prob"].shape
):
[74]:
trace = poutine.trace(model1).get_trace() trace.compute_log_prob() # optional, but allows printing of log_prob shapes print(trace.format_shapes())
Trace Shapes: Param Sites: Sample Sites: a dist | value | log_prob | b dist | 2 value | 2 log_prob | c_plate dist | value 2 | log_prob | c dist 2 | value 2 | log_prob 2 | d_plate dist | value 3 | log_prob | d dist 3 | 4 5 value 3 | 4 5 log_prob 3 | x_axis dist | value 3 | log_prob | y_axis dist | value 2 | log_prob | x dist 3 1 | value 3 1 | log_prob 3 1 | y dist 2 1 1 | value 2 1 1 | log_prob 2 1 1 | xy dist 2 3 1 | value 2 3 1 | log_prob 2 3 1 | z dist 2 3 1 | 5 value 2 3 1 | 5 log_prob 2 3 1 |
plate
内部张量子采样
plate 的主要用途之一是对数据进行子采样。因为 plate
内部数据是条件独立的, 所以子采样之后计算的损失期望估计不会变。如果要对样本数据进行子采样,您需要同时将原始数据大小和子样本大小告知 Pyro。然后 Pyro 将选择一个随机数据子集并产生一组索引。
[3]:
data = torch.arange(100.) def model2(): mean = pyro.param("mean", torch.zeros(len(data))) with pyro.plate("data", len(data), subsample_size=10) as ind: assert len(ind) == 10 # ind is a LongTensor that indexes the subsample. batch = data[ind] # Select a minibatch of data. mean_batch = mean[ind] # Take care to select the relevant per-datum parameters. # Do stuff with batch: x = pyro.sample("x", Normal(mean_batch, 1), obs=batch) assert len(x) == 10 test_model(model2, guide=lambda: None, loss=Trace_ELBO())
广播和并行计算
-
Broadcasting to allow parallel enumeration
Pyro 0.2 引入了并行枚举离散潜变量的功能。This can significantly reduce the variance of gradient estimators when learning a posterior via SVI.
To use parallel enumeration, Pyro needs to allocate tensor dimension that it can use for enumeration. To avoid conflicting with other dimensions that we want to use for plate
s, we need to declare a budget of the maximum number of tensor dimensions we’ll use. This budget is called max_plate_nesting
and is an argument to SVI (the argument is simply passed through to TraceEnum_ELBO). Usually Pyro can determine this budget on its own (it runs the (model,guide)
pair once and record what happens), but in case of dynamic model structure you may need to declare max_plate_nesting
manually.
To understand max_plate_nesting
and how Pyro allocates dimensions for enumeration, let’s revisit model1()
from above. This time we’ll map out three types of dimensions: enumeration dimensions on the left (Pyro takes control of these), batch dimensions in the middle, and event dimensions on the right.
max_plate_nesting = 3 |<--->| enumeration|batch|event -----------+-----+----- |. . .| a = sample("a", Normal(0, 1)) |. . .|2 b = sample("b", Normal(zeros(2), 1) | | .to_event(1)) | | with plate("c", 2): |. . 2| c = sample("c", Normal(zeros(2), 1)) | | with plate("d", 3): |. . 3|4 5 d = sample("d", Normal(zeros(3,4,5), 1) | | .to_event(2)) | | | | x_axis = plate("x", 3, dim=-2) | | y_axis = plate("y", 2, dim=-3) | | with x_axis: |. 3 1| x = sample("x", Normal(0, 1)) | | with y_axis: |2 1 1| y = sample("y", Normal(0, 1)) | | with x_axis, y_axis: |2 3 1| xy = sample("xy", Normal(0, 1)) |2 3 1|5 z = sample("z", Normal(0, 1).expand([5])) | | .to_event(1))
Note that it is safe to overprovision max_plate_nesting=4
but we cannot underprovision max_plate_nesting=2
(or Pyro will error). Let’s see how this works in practice.
[4]:
@config_enumerate def model3(): p = pyro.param("p", torch.arange(6.) / 6) locs = pyro.param("locs", torch.tensor([-1., 1.])) a = pyro.sample("a", Categorical(torch.ones(6) / 6)) b = pyro.sample("b", Bernoulli(p[a])) # Note this depends on a. with pyro.plate("c_plate", 4): c = pyro.sample("c", Bernoulli(0.3)) with pyro.plate("d_plate", 5): d = pyro.sample("d", Bernoulli(0.4)) e_loc = locs[d.long()].unsqueeze(-1) e_scale = torch.arange(1., 8.) e = pyro.sample("e", Normal(e_loc, e_scale) .to_event(1)) # Note this depends on d. # enumerated|batch|event dims assert a.shape == ( 6, 1, 1 ) # Six enumerated values of the Categorical. assert b.shape == ( 2, 1, 1, 1 ) # Two enumerated Bernoullis, unexpanded. assert c.shape == ( 2, 1, 1, 1, 1 ) # Only two Bernoullis, unexpanded. assert d.shape == (2, 1, 1, 1, 1, 1 ) # Only two Bernoullis, unexpanded. assert e.shape == (2, 1, 1, 1, 5, 4, 7) # This is sampled and depends on d. assert e_loc.shape == (2, 1, 1, 1, 1, 1, 1,) assert e_scale.shape == ( 7,) test_model(model3, model3, TraceEnum_ELBO(max_plate_nesting=2))
Let’s take a closer look at those dimensions. First note that Pyro allocates enumeration dims starting from the right at max_plate_nesting
: Pyro allocates dim -3 to enumerate a
, then dim -4 to enumerate b
, then dim -5 to enumerate c
, and finally dim -6 to enumerate d
. Next note that samples only have extent (size > 1) in the new enumeration dimension. This helps keep tensors small and computation cheap. (Note that the log_prob
shape will be broadcast up to contain both enumeratin shape and batch shape, so e.g. trace.nodes['d']['log_prob'].shape == (2, 1, 1, 1, 5, 4)
.)
We can draw a similar map of the tensor dimensions:
max_plate_nesting = 2 |<->| enumeration batch event ------------|---|----- 6|1 1| a = pyro.sample("a", Categorical(torch.ones(6) / 6)) 2 1|1 1| b = pyro.sample("b", Bernoulli(p[a])) | | with pyro.plate("c_plate", 4): 2 1 1|1 1| c = pyro.sample("c", Bernoulli(0.3)) | | with pyro.plate("d_plate", 5): 2 1 1 1|1 1| d = pyro.sample("d", Bernoulli(0.4)) 2 1 1 1|1 1|1 e_loc = locs[d.long()].unsqueeze(-1) | |7 e_scale = torch.arange(1., 8.) 2 1 1 1|5 4|7 e = pyro.sample("e", Normal(e_loc, e_scale) | | .to_event(1))
To automatically examine this model with enumeration semantics, we can create an enumerated trace and then use Trace.format_shapes():
[11]:
trace = poutine.trace(poutine.enum(model3, first_available_dim=-3)).get_trace() trace.compute_log_prob() # optional, but allows printing of log_prob shapes print(trace.format_shapes())
Trace Shapes: Param Sites: p 6 locs 2 Sample Sites: a dist | value 6 1 1 | log_prob 6 1 1 | b dist 6 1 1 | value 2 1 1 1 | log_prob 2 6 1 1 | c_plate dist | value 4 | log_prob | c dist 4 | value 2 1 1 1 1 | log_prob 2 1 1 1 4 | d_plate dist | value 5 | log_prob | d dist 5 4 | value 2 1 1 1 1 1 | log_prob 2 1 1 1 5 4 | e dist 2 1 1 1 5 4 | 7 value 2 1 1 1 5 4 | 7 log_prob 2 1 1 1 5 4 |
Writing parallelizable code
It can be tricky to write Pyro models that correctly handle parallelized sample sites. Two tricks help: broadcasting and ellipsis slicing. Let’s look at a contrived model to see how these work in practice. 我们的目标是写一个模型 that works both with and without enumeration.
[79]:
width, height = 8, 10 sparse_pixels = torch.LongTensor([[3, 2], [3, 5], [3, 9], [7, 1]]) enumerated = None # set to either True or False below def fun(observe): p_x = pyro.param("p_x", torch.tensor(0.1), constraint=constraints.unit_interval) p_y = pyro.param("p_y", torch.tensor(0.1), constraint=constraints.unit_interval) x_axis = pyro.plate('x_axis', width, dim=-2) y_axis = pyro.plate('y_axis', height, dim=-1) # Note that the shapes of these sites depend on whether Pyro is enumerating. with x_axis: x_active = pyro.sample("x_active", Bernoulli(p_x)) with y_axis: y_active = pyro.sample("y_active", Bernoulli(p_y)) if enumerated: assert x_active.shape == (2, 1, 1) assert y_active.shape == (2, 1, 1, 1) else: assert x_active.shape == (width, 1) assert y_active.shape == (height,) # The first trick is to broadcast. This works with or without enumeration. p = 0.1 + 0.5 * x_active * y_active if enumerated: assert p.shape == (2, 2, 1, 1) else: assert p.shape == (width, height) dense_pixels = p.new_zeros(broadcast_shape(p.shape, (width, height))) # The second trick is to index using ellipsis slicing. # This allows Pyro to add arbitrary dimensions on the left. for x, y in sparse_pixels: dense_pixels[..., x, y] = 1 if enumerated: assert dense_pixels.shape == (2, 2, width, height) else: assert dense_pixels.shape == (width, height) with x_axis, y_axis: if observe: pyro.sample("pixels", Bernoulli(p), obs=dense_pixels) def model4(): fun(observe=True) def guide4(): fun(observe=False) # Test without enumeration. enumerated = False test_model(model4, guide4, Trace_ELBO()) # Test with enumeration. enumerated = True test_model(model4, config_enumerate(guide4, "parallel"), TraceEnum_ELBO(max_plate_nesting=2))
plate
内部自动广播
在 Pyro 中,`pyro.plate` 用于声明一组条件独立的变量,并且它可以帮助自动扩展样本的维度,以满足 `pyro.sample` 语句中关于 `batch_shape` 的要求。这种自动扩展与手动编写 `.expand()` 语句来广播分发的批次形状是等价的。
以下是对您提供信息的解释:
- **自动扩展样本维度**:在定义模型(`model`)或引导(`guide`)时,Pyro 的 `pyro.plate` 可以自动扩展样本的维度,以确保它们与 `pyro.sample` 调用中的 `batch_shape` 兼容。这意味着,如果在一个 `plate` 的上下文中进行采样,Pyro 会自动调整样本的数量,以匹配 `plate` 指定的大小。
- **与 `.expand()` 等价的广播**:虽然 `pyro.plate` 提供了一种方便的方式来自动扩展样本维度,但它本质上与手动在分布上使用 `.expand()` 方法来指定批次大小是相同的。`.expand()` 方法用于从分布中抽取多个独立同分布的样本,而 `pyro.plate` 则在概率编程的上下文中提供了一种更自然的方式来声明这种独立性。
- **示例演示**:您提到将使用之前部分中的 `model4` 来演示这一点。在这种情况下,`model4` 可能会展示如何在 `pyro.plate` 的上下文中使用 `pyro.sample`,以及如何通过 `plate` 自动扩展样本维度,或者通过手动使用 `.expand()` 来达到相同的目的。
- **对先前代码的更改**:您提到了对先前代码的更改,但具体更改的内容没有提供。通常,这些更改可能涉及到使用 `pyro.plate` 替代或补充 `.expand()` 语句,以改进模型定义的清晰度和模块化。
理解这一点对于有效地使用 Pyro 构建概率模型非常重要,因为它允许您以一种声明性的方式编写代码,同时保持对底层概率分布的精确控制。如果您需要进一步的具体示例或解释,请提供更多的上下文或代码示例。
所有的 model/guide 定义中,我们依赖于 pyro.plate 自动扩展样本维度,来满足 pyro.sample
语句中关于 batch shape 的约束。However this broadcasting is equivalent to hand-annotated .expand()
statements.
我们将使用 previous section 中的 model4
对此进行演示。 请注意对先前代码的以下更改:
-
For the purpose of this example, we will only consider “parallel” enumeration, but broadcasting should work as expected without enumeration or with “sequential” enumeration.
-
We have separated out the sampling function which returns the tensors corresponding to the active pixels. Modularizing the model code into components is a common practice, and helps with maintainability of large models.
-
We would also like to use the
pyro.plate
construct to parallelize the ELBO estimator over num_particles. This is done by wrapping the contents of model/guide inside an outermostpyro.plate
context.
[78]:
num_particles = 100 # Number of samples for the ELBO estimator width = 8 height = 10 sparse_pixels = torch.LongTensor([[3, 2], [3, 5], [3, 9], [7, 1]]) def sample_pixel_locations_no_broadcasting(p_x, p_y, x_axis, y_axis): with x_axis: x_active = pyro.sample("x_active", Bernoulli(p_x).expand([num_particles, width, 1])) with y_axis: y_active = pyro.sample("y_active", Bernoulli(p_y).expand([num_particles, 1, height])) return x_active, y_active def sample_pixel_locations_full_broadcasting(p_x, p_y, x_axis, y_axis): with x_axis: x_active = pyro.sample("x_active", Bernoulli(p_x)) with y_axis: y_active = pyro.sample("y_active", Bernoulli(p_y)) return x_active, y_active def sample_pixel_locations_partial_broadcasting(p_x, p_y, x_axis, y_axis): with x_axis: x_active = pyro.sample("x_active", Bernoulli(p_x).expand([width, 1])) with y_axis: y_active = pyro.sample("y_active", Bernoulli(p_y).expand([height])) return x_active, y_active def fun(observe, sample_fn): p_x = pyro.param("p_x", torch.tensor(0.1), constraint=constraints.unit_interval) p_y = pyro.param("p_y", torch.tensor(0.1), constraint=constraints.unit_interval) x_axis = pyro.plate('x_axis', width, dim=-2) y_axis = pyro.plate('y_axis', height, dim=-1) with pyro.plate("num_particles", 100, dim=-3): x_active, y_active = sample_fn(p_x, p_y, x_axis, y_axis) # Indices corresponding to "parallel" enumeration are appended # to the left of the "num_particles" plate dim. assert x_active.shape == (2, 1, 1, 1) assert y_active.shape == (2, 1, 1, 1, 1) p = 0.1 + 0.5 * x_active * y_active assert p.shape == (2, 2, 1, 1, 1) dense_pixels = p.new_zeros(broadcast_shape(p.shape, (width, height))) for x, y in sparse_pixels: dense_pixels[..., x, y] = 1 assert dense_pixels.shape == (2, 2, 1, width, height) with x_axis, y_axis: if observe: pyro.sample("pixels", Bernoulli(p), obs=dense_pixels) def test_model_with_sample_fn(sample_fn): def model(): fun(observe=True, sample_fn=sample_fn) @config_enumerate def guide(): fun(observe=False, sample_fn=sample_fn) test_model(model, guide, TraceEnum_ELBO(max_plate_nesting=3)) test_model_with_sample_fn(sample_pixel_locations_no_broadcasting) test_model_with_sample_fn(sample_pixel_locations_full_broadcasting) test_model_with_sa mple_fn(sample_pixel_locations_partial_broadcasting)
在 Pyro 中,`pyro.plate` 是一个强大的工具,用于声明条件独立性的维度,并且它可以极大地简化并行化操作和代码模块化。以下是对您提供的段落的理解:
1. **手动记录和扩展**:在第一个采样函数中,需要手动记录并扩展 `Bernoulli` 分布的批量形状(`batch_shape`),以适应由 `pyro.plate` 上下文添加的条件独立维度。这意味着,采样函数 `sample_pixel_locations` 需要知道 `num_particles`、`width` 和 `height` 这些全局变量的值,这并不理想,因为它依赖于全局作用域中的变量。
2. **`pyro.plate` 的可选 `size` 参数**:`pyro.plate` 的第二个参数是可选的 `size` 参数,它需要提供以进行隐式广播,以便它可以推断每个采样点的批量形状需求。这有助于 Pyro 理解在 `pyro.plate` 上下文中所需的批量大小。
3. **现有的 `batch_shape` 必须可广播**:在特定的例子中,`Bernoulli(p_x)` 有一个空的批量形状,这意味着它可以与任何尺寸广播,是通用的。
4. **通过 `pyro.plate` 实现并行化**:使用 `pyro.plate` 可以通过张量化操作轻松实现并行化。`pyro.plate` 通过允许模型组件在不知道它们后续可能被嵌入的 `plate` 上下文的情况下编写,从而帮助代码模块化。
简而言之,`pyro.plate` 提供了一种声明多个条件独立变量的方法,这些变量可以并行处理,而不需要显式地编写循环或迭代。它使得模型定义更加简洁,并且模型组件可以独立于它们被嵌入的 `plate` 上下文进行编写。这样,您可以构建更灵活、更易于维护的概率模型,同时利用 Pyro 的并行化能力来提高计算效率。
In the first sampling function, we had to do some manual book-keeping and expand the Bernoulli
distribution’s batch shape to account for the conditionally independent dimensions added by the pyro.plate
contexts. In particular, note how sample_pixel_locations
needs knowledge of num_particles
, width
and height
and is accessing these variables from the global scope, which is not ideal.
-
The second argument to
pyro.plate
, i.e. the optionalsize
argument needs to be provided for implicit broadasting, so that it can infer the batch shape requirement for each of the sample sites. -
The existing
batch_shape
of the sample site must be broadcastable with the size of thepyro.plate
contexts. In our particular example,Bernoulli(p_x)
has an empty batch shape which is universally broadcastable.
Note how simple it is to achieve parallelization via tensorized operations using pyro.plate
! pyro.plate
also helps in code modularization because model components can be written agnostic of the plate
contexts in which they may subsequently get embedded in.