Pyro 入门

概率是处理不确定性推理的数学，正如微积分是处理变化率推理的数学一样。它为理解现代机器学习和人工智能的大部分内容提供了一个统一的理论框架：用概率语言构建的模型可以捕捉复杂的推理，了解它们不知道的东西，并在没有监督的情况下揭示数据中的结构。

直接指定概率模型可能会很麻烦，实现它们也可能非常容易出错。概率编程语言 (PPL) 通过将概率与编程语言的表达能力相结合来解决这些问题。概率程序是普通确定性计算和随机采样值的混合，代表数据的生成过程。

通过观察概率程序的输出，我们可以描述一个推断问题，大致可以翻译为：“如果这个随机选择有一个特定的观测值，那么什么必须是真实的？”PPL 明确强制了概率数学中模型规范、要回答的查询以及计算答案的算法之间的关注点分离。

Pyro 是一个基于 Python 和 PyTorch 构建的概率编程语言。Pyro 程序就是 Python 程序，其主要推断技术是随机变分推断，它将抽象的概率计算转换为使用 PyTorch 中的随机梯度下降解决的具体优化问题，使得概率方法可以应用于以前难以处理的模型和数据集规模。

在本教程中，我们将通过一个涉及线性回归的数据分析问题示例，对概率机器学习和 Pyro 概率编程的基本概念进行简要的、带有观点的介绍。线性回归是机器学习中最常见、最基本的任务之一。我们将看到如何使用 Pyro 的建模语言和推断算法将不确定性纳入回归系数的估计中。

大纲

设置

让我们首先导入所需的模块。

[41]:

%reset -s -f

[42]:

import logging
import os

import torch
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import pyro

[43]:

smoke_test = ('CI' in os.environ)
assert pyro.__version__.startswith('1.9.1')

pyro.enable_validation(True)
pyro.set_rng_seed(1)
logging.basicConfig(format='%(message)s', level=logging.INFO)

# Set matplotlib settings
%matplotlib inline
plt.style.use('default')

背景：概率机器学习

大多数数据分析问题都可以理解为对三个基本高级问题的阐述

  1. 在观察任何数据之前，我们对问题了解多少？

  2. 根据先验知识，我们可以从数据中得出什么结论？

  3. 这些结论有意义吗？

在数据科学和机器学习的概率或贝叶斯方法中，我们以概率分布上的数学运算的形式来形式化这些问题。

背景：概率模型

首先，我们将对问题中变量及其之间关系的了解全部表达出来，形成一个概率模型，或一组随机变量的联合概率分布。模型具有观测值 \({\bf x}\) 和潜随机变量 \({\bf z}\) 以及参数 \(\theta\)。它通常具有以下形式的联合密度函数

\[p_{\theta}({\bf x}, {\bf z}) = p_{\theta}({\bf x}|{\bf z}) p_{\theta}({\bf z})\]

公式中潜变量上的分布 \(p_{\theta}({\bf z})\) 称为先验，给定潜变量的观测变量分布 \(p_{\theta}({\bf x}|{\bf z})\) 称为似然。

我们通常要求构成模型 \(p_{\theta}({\bf x}, {\bf z})\) 的各种条件概率分布 \(p_i\) 具有以下属性（Pyro 和 PyTorch Distributions 中提供的分布通常满足这些属性）

  • 我们可以有效地从每个 \(p_i\) 中采样

  • 我们可以有效地计算逐点概率密度 \(p_i\)

  • \(p_i\) 对参数 \(\theta\) 可微分

概率模型通常用标准图形表示法描绘，用于可视化和交流，总结如下，尽管在 Pyro 中可以表示没有固定图形结构的模型。在有很多重复的模型中，使用板 (plate) 表示法很方便，之所以这样称呼是因为它在图形上显示为围绕变量的矩形“板”，表示内部随机变量的多个独立副本。

Screen%20Shot%202021-12-13%20at%201.31.34%20PM.png

背景：推断、学习和评估

一旦我们指定了模型，贝叶斯法则告诉我们如何使用它来执行推断，或者从数据中得出关于潜变量的结论，通过计算 \(\bf z\) 上的后验分布

\[ p_{\theta}({\bf z} | {\bf x}) = \frac{p_{\theta}({\bf x} , {\bf z})}{ \int \! d{\bf z}\; p_{\theta}({\bf x} , {\bf z}) }\]

为了检查建模和推断的结果，我们想知道模型对观测数据 \(x\) 的拟合程度，我们可以用证据或边际似然来量化

\[p_{\theta}({\bf x}) = \int \! d{\bf z}\; p_{\theta}({\bf x} , {\bf z})\]

还可以对新数据进行预测，我们可以使用后验预测分布来实现

\[p_{\theta}(x' | {\bf x}) = \int \! d{\bf z}\; p_{\theta}(x' | {\bf z}) p_{\theta}({\bf z} | {\bf x})\]

最后，通常需要从观测数据 \(x\) 中学习模型的参数 \(\theta\)，我们可以通过最大化边际似然来实现

\[\theta_{\rm{max}} = \rm{argmax}_\theta p_{\theta}({\bf x}) = \rm{argmax}_\theta \int \! d{\bf z}\; p_{\theta}({\bf x} , {\bf z})\]

示例：地理与国民收入

以下示例改编自 Richard McElreath 的优秀书籍 Statistical Rethinking 的第 7 章，建议读者回顾该书，以便对更广泛的贝叶斯数据分析实践有一个易于理解的介绍（所有章节的 Pyro 代码均可获取）。

我们希望探索一个国家的地形崎岖度指数（数据集中变量rugged）与其人均 GDP 之间的关系。特别是，原论文作者（“Ruggedness: The blessing of bad geography in Africa”）指出，在非洲以外地区，地形崎岖或恶劣的地理条件与较差的经济表现有关，但对于非洲国家而言，崎岖地形对收入产生了相反的影响。

让我们查看数据并研究这种关系。我们将重点关注数据集中的三个特征

  • rugged：量化地形崎岖度指数；

  • cont_africa：给定国家是否在非洲；

  • rgdppc_2000：2000 年的实际人均 GDP；

[44]:

DATA_URL = "https://d2hg8soec8ck9v.cloudfront.net/datasets/rugged_data.csv"
data = pd.read_csv(DATA_URL, encoding="ISO-8859-1")
df = data[["cont_africa", "rugged", "rgdppc_2000"]]

响应变量 GDP 高度偏斜，因此在继续之前我们将对其进行对数变换。

[45]:

df = df[np.isfinite(df.rgdppc_2000)]
df["rgdppc_2000"] = np.log(df["rgdppc_2000"])

然后我们将该数据框背后的 Numpy 数组转换为 torch.Tensor，以便使用 PyTorch 和 Pyro 进行分析。

[46]:

train = torch.tensor(df.values, dtype=torch.float)
is_cont_africa, ruggedness, log_gdp = train[:, 0], train[:, 1], train[:, 2]

数据可视化表明，崎岖度和 GDP 之间确实可能存在关系，但需要进一步分析来确认。我们将看到如何通过贝叶斯线性回归在 Pyro 中实现这一点。

[47]:

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6), sharey=True)
african_nations = df[df["cont_africa"] == 1]
non_african_nations = df[df["cont_africa"] == 0]
sns.scatterplot(x=non_african_nations["rugged"],
                y=non_african_nations["rgdppc_2000"],
                ax=ax[0])
ax[0].set(xlabel="Terrain Ruggedness Index",
          ylabel="log GDP (2000)",
          title="Non African Nations")
sns.scatterplot(x=african_nations["rugged"],
                y=african_nations["rgdppc_2000"],
                ax=ax[1])
ax[1].set(xlabel="Terrain Ruggedness Index",
          ylabel="log GDP (2000)",
          title="African Nations");
_images/intro_long_20_0.png

Pyro 中的模型

Pyro 中的概率模型被指定为 Python 函数 model(*args, **kwargs)，这些函数使用特殊的原语函数从潜变量生成观测数据，这些原语函数的行为可以根据正在执行的高级计算由 Pyro 内部机制更改。

具体来说，model() 的不同数学部分通过以下映射进行编码

  1. 潜随机变量 \(z\) \(\Longleftrightarrow\) pyro.sample

  2. 观测随机变量 \(x\) \(\Longleftrightarrow\) pyro.sample 并带有 obs 关键字参数

  3. 可学习参数 \(\theta\) \(\Longleftrightarrow\) pyro.param

  4. 板 \(\Longleftrightarrow\) pyro.plate 上下文管理器

下面我们将在线性回归的第一个 Pyro 模型上下文中详细检查这些组件。

示例模型：最大似然线性回归

如果我们将线性回归预测器 \(\beta X + \alpha\) 的公式写成 Python 表达式，我们得到以下结果

mean = a + b_a * is_cont_africa + b_r * ruggedness + b_ar * is_cont_africa * ruggedness

为了将其构建成一个完整的概率模型来处理我们的数据集，我们需要将回归系数 \(\alpha\) 和 \(\beta\) 设置为可学习参数，并在预测均值周围添加观测噪声。我们可以使用上面介绍的 Pyro 原语来表达这一点，并使用 pyro.render_model() 可视化结果模型

[48]:

import pyro.distributions as dist
import pyro.distributions.constraints as constraints

def simple_model(is_cont_africa, ruggedness, log_gdp=None):
    a = pyro.param("a", lambda: torch.randn(()))
    b_a = pyro.param("bA", lambda: torch.randn(()))
    b_r = pyro.param("bR", lambda: torch.randn(()))
    b_ar = pyro.param("bAR", lambda: torch.randn(()))
    sigma = pyro.param("sigma", lambda: torch.ones(()), constraint=constraints.positive)

    mean = a + b_a * is_cont_africa + b_r * ruggedness + b_ar * is_cont_africa * ruggedness

    with pyro.plate("data", len(ruggedness)):
        return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

pyro.render_model(simple_model, model_args=(is_cont_africa, ruggedness, log_gdp), render_distributions=True)

[48]:
_images/intro_long_24_0.svg

上面的图不显示模型参数 a, b_a, b_r, b_ar, 和 sigma。我们可以设置 render_params=True 来渲染模型参数。

[49]:

pyro.render_model(simple_model, model_args=(is_cont_africa, ruggedness, log_gdp), render_distributions=True, render_params=True)

[49]:
_images/intro_long_26_0.svg

参数 render_distributions = True 将显示参数的约束。例如，sigma 是一个标准差，它应该是非负的。因此，它的约束显示为“sigma \(\in\) GreaterThan(lower_bound=0.0)”。

学习 simple_model 的参数构成了最大似然估计，并产生回归系数的点估计。然而，在这个例子中，我们的数据可视化表明我们不应该对回归系数的任何单一值过于自信。相比之下，完全贝叶斯方法将对不同的可能参数值以及模型的预测产生不确定性估计。

在构建贝叶斯版本的线性模型之前，让我们停下来仔细看看这段 Pyro 代码。

背景：pyro.sample 原语

Pyro 中的概率程序是围绕来自原始概率分布的样本构建的，用 pyro.sample 标记

def sample(
    name: str,
    fn: pyro.distributions.Distribution,
    *,
    obs: typing.Optional[torch.Tensor] = None,
    infer: typing.Optional[dict] = None
) -> torch.Tensor:
    ...

在上面的模型 simple_model 中，这一行

return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

可以表示一个潜变量或一个观测变量，具体取决于是否给 simple_model 提供了 log_gdp 值。当不提供 log_gdp 且 "obs" 是潜变量时，它等价于（暂时忽略 pyro.plate 上下文，假设 len(ruggedness) == 1）调用分布底层 .sample 方法

return dist.Normal(mean, sigma).sample()

这种解释是偶尔将 Pyro 程序称为随机函数的原因，这是一个在 Pyro 早期文档中使用的一个相当模糊的术语。

当 simple_model 接收到 log_gdp 参数且 "obs" 是观测值时，pyro.sample 语句

return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

将始终返回 log_gdp

return log_gdp

然而，请注意，**当任何样本语句被观测到时，模型中所有其他样本语句的累积效应都会根据贝叶斯法则发生改变**；Pyro 的推断算法的工作就是“反向运行程序”，并为模型中的所有pyro.sample语句分配数学上一致的值。

这时一个合理的问题是，为什么 pyro.sample 和其他原语必须有名称。名称被用户和 Pyro 的内部机制用来分离模型、观测值和推断算法的规范，这是概率编程语言的一个关键卖点。要看到一个例子，我们可以看一下更高阶原语 pyro.condition，它解决了对单个 Pyro 模型编写多个查询的问题

def condition(
    model: Callable[..., T],
    data: Dict[str, torch.Tensor]
) -> Callable[..., T]:
    ...

pyro.condition 接受一个模型和一个（可能为空的）字典，该字典将名称映射到观测值，并将每个观测值传递给其名称指示的 pyro.sample 语句。在我们的示例 simple_model 的上下文中，我们可以将 log_gdp 从参数中移除，并将其替换为更简单的接口

def simpler_model(is_cont_africa, ruggedness): ...

conditioned_model = pyro.condition(simpler_model, data={"obs": log_gdp})

其中 conditioned_model 等价于

conditioned_model = functools.partial(simple_model, log_gdp=log_gdp)

背景：pyro.param 原语

我们模型中使用的下一个原语 pyro.param 是一个用于从 Pyro 的键值参数存储中读取和写入的前端

def param(
    name: str,
    init: Optional[Union[torch.Tensor, Callable[..., torch.Tensor]]] = None,
    *,
    constraint: torch.distributions.constraints.Constraint = constraints.real
) -> torch.Tensor:
    ...

与 pyro.sample 类似，pyro.param 总是以名称作为其第一个参数调用。第一次使用特定名称调用 pyro.param 时，它会将第二个参数 init 指定的初始值存储在参数存储中，然后返回该值。此后，当使用该名称调用时，它将返回参数存储中的值，而忽略任何其他参数。参数初始化后，不再需要指定 init 来检索其值（例如 pyro.param("a")）。

第二个参数 init 可以是 torch.Tensor 或不带参数并返回张量的函数。后一种形式很有用，因为它可以避免重复构建只在模型第一次运行时使用的初始值。

与 PyTorch 的 torch.nn.Parameter 不同，Pyro 中的参数可以明确约束在 \(\mathbb{R}^n\) 的各种子集上，这是一个重要特性，因为许多基本概率分布的参数域受到限制。例如，Normal 分布的 scale 参数必须为正。传递给 pyro.param 的可选第三个参数 constraint 是一个在参数初始化时存储的 torch.distributions.constraints.Constraint 对象；约束在每次更新后会重新应用。Pyro 内置了大量预定义的约束。

pyro.param 的值在模型调用之间保持不变，除非参数存储被优化算法更新或通过 pyro.clear_param_store() 清除。与 pyro.sample 不同，pyro.param 可以在模型中多次使用相同的名称调用；每次使用相同的名称调用都会返回相同的值。全局参数存储本身可以通过调用 pyro.get_param_store() 访问。

在上面的模型 simple_model 中，语句

a = pyro.param("a", lambda: torch.randn(()))

在概念上类似于下面的代码，但附加了一些跟踪、序列化和约束管理功能。

simple_param_store = {}
...
def simple_model():
    a_init = lambda: torch.randn(())
    a = simple_param_store["a"] if "a" in simple_param_store else a_init()
    ...

虽然本入门教程使用 pyro.param 进行参数管理，但 Pyro 也兼容 PyTorch 熟悉的 torch.nn.Module API，通过 pyro.nn.PyroModule 实现。

背景：pyro.plate 原语

pyro.plate 是 Pyro 对板表示法的形式化编码，该表示法在概率机器学习中广泛使用，用于简化具有大量条件独立同分布随机变量的模型的可视化和分析。

def plate(
    name: str,
    size: int,
    *,
    dim: Optional[int] = None,
    **other_kwargs
) -> contextlib.AbstractContextManager:
    ...

概念上，pyro.plate 语句等价于一个 for 循环。在 simple_model 中，我们可以用 Python for 循环替换以下几行

with pyro.plate("data", len(ruggedness)):
    return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

用 Python 的 for 循环

result = torch.empty_like(ruggedness)
for i in range(len(ruggedness)):
    result[i] = pyro.sample(f"obs_{i}", dist.Normal(mean, sigma), obs=log_gdp[i] if log_gdp is not None else None)
return result

当重复变量的数量（本例中为 len(ruggedness)）很大时，使用 Python 循环会相当慢。由于循环中的每次迭代都相互独立，因此 pyro.plate 使用 PyTorch 的数组广播在一个向量化操作中并行执行迭代，如下面的等效向量化代码所示

mean = mean.unsqueeze(-1).expand((len(ruggedness),))
sigma = sigma.unsqueeze(-1).expand((len(ruggedness),))
return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

实际操作中，在编写 Pyro 程序时，pyro.plate 主要用作向量化工具。然而，如SVI 教程第二部分所述，它也用于管理数据子采样，并作为推断算法的信号，表明某些变量是独立的。

示例：从最大似然回归到贝叶斯回归

为了使我们的线性回归模型成为贝叶斯模型，我们需要对参数 \(\alpha \in \mathbb{R}\) 和 \(\beta \in \mathbb{R}^3\) (此处展开为标量 b_a, b_r, 和 b_ar) 指定先验分布。这些是概率分布，代表我们在观察任何数据之前对 \(\alpha\) 和 \(\beta\) 合理值的信念。我们还将添加一个随机尺度参数 \(\sigma\) 来控制观测噪声。

在 Pyro 中表达这个贝叶斯线性回归模型非常直观：我们只需将每个 pyro.param 语句替换为带有描述每个参数先验信念的Pyro Distribution 对象的 pyro.sample 语句。

对于常数项 \(\alpha\)，我们使用标准差较大的 Normal 先验，表示我们对基线 GDP 的先验知识相对较少。对于其他回归系数，我们使用标准的 Normal 先验（均值在 0），表示我们对协变量与 GDP 之间关系是正向还是负向缺乏先验知识。对于观测噪声 \(\sigma\)，我们使用一个以 0 为下界的平坦先验，因为该值必须为正才能成为有效的标准差。

[50]:

def model(is_cont_africa, ruggedness, log_gdp=None):
    a = pyro.sample("a", dist.Normal(0., 10.))
    b_a = pyro.sample("bA", dist.Normal(0., 1.))
    b_r = pyro.sample("bR", dist.Normal(0., 1.))
    b_ar = pyro.sample("bAR", dist.Normal(0., 1.))
    sigma = pyro.sample("sigma", dist.Uniform(0., 10.))

    mean = a + b_a * is_cont_africa + b_r * ruggedness + b_ar * is_cont_africa * ruggedness

    with pyro.plate("data", len(ruggedness)):
        return pyro.sample("obs", dist.Normal(mean, sigma), obs=log_gdp)

pyro.render_model(model, model_args=(is_cont_africa, ruggedness, log_gdp), render_distributions=True)

[50]:
_images/intro_long_34_0.svg

Pyro 中的推断

背景：变分推断

介绍中的所有计算（后验分布、边际似然和后验预测分布）都需要执行积分，而这些积分通常是不可能或计算上难以处理的。

虽然 Pyro 支持许多不同的精确和近似推断算法，但支持最好的是变分推断，它提供了一个统一的方案来寻找 \(\theta_{\rm{max}}\) 并计算真实未知后验分布 \(p_{\theta_{\rm{max}}}({\bf z} | {\bf x})\) 的可处理近似 \(q_{\phi}({\bf z})\)，通过将难以处理的积分转换为 \(p\) 和 \(q\) 的泛函优化。下图概念性地描述了这个过程，更全面的数学介绍可在SVI 教程第一部分中找到。

大多数概率分布（下图中浅色椭圆），尤其是那些对应于贝叶斯后验分布的，过于复杂，无法直接表示，因此我们必须定义一个较小的子空间，由实值参数 \(\phi\) 索引，其中分布 \(q_{\phi}({\bf z})\) 在构造上保证易于采样（下图中深色圆圈），但可能不包含真实的后验分布 \(p_{\theta}({\bf z} | {\bf x})\)（下图中红色星形）。

变分推断通过搜索变分分布空间，找到一个根据某种距离或散度度量（下图中黑色箭头）与真实后验分布最相似的分布（下图中黄色星形）来近似真实后验分布。

Screen%20Shot%202021-12-10%20at%203.13.09%20PM.png

然而，衡量概率分布之间距离或散度的方法有很多种。我们应该选择哪一种呢？如图所示，一个理论上吸引人的选择是 Kullback-Leibler 散度 \(KL(q_{\phi}({\bf z}) || p_{\theta}({\bf z} | {\bf x}))\)，但直接计算它需要预先知道真实的后验分布，这会适得其反。

更重要的是，我们对优化这种散度感兴趣，这听起来可能更难，但实际上可以使用贝叶斯定理将 \(KL(q_{\phi}({\bf z}) || p_{\theta}({\bf z} | {\bf x}))\) 的定义重写为一个不依赖于 \(q_{\phi}\) 的难以处理的常数与一个可处理项（称为证据下界 (ELBO)，定义如下）之间的差值。因此，最大化这个可处理项将产生与最小化原始 KL 散度相同的解。

背景：作为灵活近似后验分布的“guide”程序

在变分推断中，我们引入一个参数化分布 \(q_{\phi}({\bf z})\) 来近似真实的后验分布，其中 \(\phi\) 称为变分参数。在许多文献中，这个分布称为变分分布，而在 Pyro 中，它被称为 guide（少了一个音节！）。

就像模型一样，guide 被编码为一个 Python 程序 guide()，其中包含 pyro.sample 和 pyro.param 语句。它不包含观测数据，因为 guide 需要是 properly normalized 的分布，以便易于从中采样。请注意，Pyro 强制要求 model() 和 guide() 采用相同的参数。允许 guide 是任意的 Pyro 程序开辟了编写 guide 族的可能性，这些 guide 族可以捕捉真实后验分布中更多特定于问题的结构，从而仅在有用的方向上扩展搜索空间，如下图所示。

Screen%20Shot%202021-12-10%20at%203.13.22%20PM.png

变分推断的数学对 guides 有什么限制？由于 guide 是后验分布 \(p_{\theta_{\rm{max}}}({\bf z} | {\bf x})\) 的近似，guide 需要为模型中的所有潜随机变量提供有效的联合概率密度。回想一下，当在 Pyro 中使用原语语句 pyro.sample() 指定随机变量时，第一个参数表示随机变量的名称。这些名称将用于对齐模型和 guide 中的随机变量。更明确地说，如果模型包含随机变量 z_1

def model():
    pyro.sample("z_1", ...)

则 guide 需要有一个对应的 sample 语句

def guide():
    pyro.sample("z_1", ...)

则 guide 需要有一个匹配的 sample 语句

尽管它提供了灵活性，但手动编写 guide 可能困难且乏味，特别是对于新用户而言。在可能的情况下，我们推荐使用 autoguides，它们是用于自动生成 Pyro 提供的模型中常见的 guide 系列的“配方”，位于 pyro.infer.autoguide 模块中。下一节将演示这两种方法。

在两种情况下使用的分布可以不同，但名称必须一对一地对齐。

尽管提供了灵活性，手动编写 guides 可能很困难且繁琐，尤其是对于新用户。在可能的情况下，我们建议使用autoguides，它是 Pyro 中包含的、用于从给定模型自动生成常见 guide 族的模板，位于pyro.infer.autoguide中。下一节将演示这两种方法。

示例：Pyro 中贝叶斯线性回归的平均场变分近似

[51]:

def custom_guide(is_cont_africa, ruggedness, log_gdp=None):
    a_loc = pyro.param('a_loc', lambda: torch.tensor(0.))
    a_scale = pyro.param('a_scale', lambda: torch.tensor(1.),
                         constraint=constraints.positive)
    sigma_loc = pyro.param('sigma_loc', lambda: torch.tensor(0.))
    weights_loc = pyro.param('weights_loc', lambda: torch.randn(3))
    weights_scale = pyro.param('weights_scale', lambda: torch.ones(3),
                               constraint=constraints.positive)
    a = pyro.sample("a", dist.Normal(a_loc, a_scale))
    b_a = pyro.sample("bA", dist.Normal(weights_loc[0], weights_scale[0]))
    b_r = pyro.sample("bR", dist.Normal(weights_loc[1], weights_scale[1]))
    b_ar = pyro.sample("bAR", dist.Normal(weights_loc[2], weights_scale[2]))
    sigma = pyro.sample("sigma", dist.LogNormal(sigma_loc, torch.tensor(0.05)))  # fixed scale for simplicity
    return {"a": a, "b_a": b_a, "b_r": b_r, "b_ar": b_ar, "sigma": sigma}

对于我们正在进行的贝叶斯线性回归示例，我们将使用一个将模型中未观测参数的分布建模为具有对角协方差的高斯分布的 guide，即它假设潜变量之间没有相关性（正如我们将看到的那样，这是一个很强的假设）。这被称为平均场近似，这是一个从物理学中借来的术语，这种近似最初就是在那里发明的。

[52]:

pyro.render_model(custom_guide, model_args=(is_cont_africa, ruggedness, log_gdp), render_params=True)

[52]:
_images/intro_long_44_0.svg

为了完整起见，我们首先手动编写一个这种形式的 guide 程序。

我们可以使用 pyro.render_model 来可视化 custom_guide，这表明随机变量确实是相互独立的，图之间没有边缘连接。

[53]:

auto_guide = pyro.infer.autoguide.AutoNormal(model)

Pyro 还包含大量“autoguides”集合，它们可以从给定的模型中自动生成 guide 程序。与我们手写的 guide 一样，所有 pyro.autoguide.AutoGuide 实例（它们本身只是接受与模型相同参数的函数）都为其包含的每个 pyro.sample 站点返回一个值字典。

最简单的 autoguide 类是 AutoNormal，它可以通过一行代码自动生成一个与我们上面手写的 guide 等效的 guide

guide 本身并不能完全指定推断算法：它只是描述了由参数索引的可能近似后验分布的搜索空间（上图中深色圆圈）以及由初始参数值确定的空间中的初始点。然后，我们必须通过解决参数的优化问题（上图中黄色星形），将该初始分布向真实后验分布（上图中红色星形）移动。如何制定和解决这个优化问题是接下来两节的主题。

背景：估计和优化证据下界 (ELBO)

我们将要优化的模型 \(p_{\theta}({\bf x}, {\bf z})\) 和 guide \(q_{\phi}({\bf z})\) 的泛函是 ELBO，它被定义为关于从 guide 采样得到的期望

\[{\rm ELBO} \equiv \mathbb{E}_{q_{\phi}({\bf z})} \left [ \log p_{\theta}({\bf x}, {\bf z}) - \log q_{\phi}({\bf z}) \right]\]

根据假设，我们可以计算期望内部的所有概率，并且由于 guide \(q\) 被假定为可采样的参数分布，因此我们可以计算该数量的蒙特卡洛估计以及相对于模型和 guide 参数 \(\theta,\phi\) 的梯度 \(\nabla_{\theta,\phi}ELBO\)。

使用这些梯度估计通过随机梯度下降优化 ELBO 在模型和 guide 参数 \(\theta,\phi\) 上的过程有时被称为随机变分推断 (SVI)；有关 SVI 的扩展介绍，请参阅SVI 第一部分。

示例：通过随机变分推断 (SVI) 进行贝叶斯回归

[54]:

adam = pyro.optim.Adam({"lr": 0.02})
elbo = pyro.infer.Trace_ELBO()
svi = pyro.infer.SVI(model, auto_guide, adam, elbo)

Pyro 包含许多不同的实现，用于估计 ELBO（在上一节中进行了数学定义），每种实现都以略微不同的方式计算损失和梯度，并具有不同的权衡。在本教程中，我们将仅使用 pyro.infer.Trace_ELBO，这样做总是正确且安全的；其他 ELBO 估计器可能为某些模型和 guide 提供计算或统计优势。

[55]:

%%time
pyro.clear_param_store()

# These should be reset each training loop.
auto_guide = pyro.infer.autoguide.AutoNormal(model)
adam = pyro.optim.Adam({"lr": 0.02})  # Consider decreasing learning rate.
elbo = pyro.infer.Trace_ELBO()
svi = pyro.infer.SVI(model, auto_guide, adam, elbo)

losses = []
for step in range(1000 if not smoke_test else 2):  # Consider running for more steps.
    loss = svi.step(is_cont_africa, ruggedness, log_gdp)
    losses.append(loss)
    if step % 100 == 0:
        logging.info("Elbo loss: {}".format(loss))

plt.figure(figsize=(5, 2))
plt.plot(losses)
plt.xlabel("SVI step")
plt.ylabel("ELBO loss");

Elbo loss: 694.9404826164246
Elbo loss: 524.3822101354599
Elbo loss: 475.66820669174194
Elbo loss: 399.99088364839554
Elbo loss: 315.23274326324463
Elbo loss: 254.76771265268326
Elbo loss: 248.237040579319
Elbo loss: 248.42670530080795
Elbo loss: 248.46450632810593
Elbo loss: 257.41463351249695

CPU times: user 6.47 s, sys: 241 µs, total: 6.47 s
Wall time: 6.28 s

[55]:

Text(0, 0.5, 'ELBO loss')
_images/intro_long_52_3.png

我们将在示例模型中使用 SVI 进行推断，演示 Pyro 如何使用 PyTorch 的随机梯度下降实现来优化 pyro.infer.Trace_ELBO 对象的输出，我们将其传递给 pyro.infer.SVI，这是一个辅助类，其 step() 方法负责计算损失和参数梯度，以及对参数应用更新和约束。

这里的 pyro.optim.Adam 是 PyTorch 优化器 torch.optim.Adam 的一个薄封装（有关讨论，请参阅此处）。pyro.optim 中的优化器用于优化和更新 Pyro 参数存储中的参数值。特别需要注意的是，我们不需要将可学习参数传递给优化器，因为这由 guide 代码确定，并在 SVI 类中自动进行。要执行 ELBO 梯度步长，我们只需调用 SVI 的 step 方法。我们传递给 SVI.step 的数据参数将传递给 model() 和 guide()。完整的训练循环如下所示

[56]:

for name, value in pyro.get_param_store().items():
    print(name, pyro.param(name).data.cpu().numpy())

AutoNormal.locs.a 9.173145
AutoNormal.scales.a 0.0703669
AutoNormal.locs.bA -1.8474661
AutoNormal.scales.bA 0.1407009
AutoNormal.locs.bR -0.19032118
AutoNormal.scales.bR 0.044044234
AutoNormal.locs.bAR 0.35599768
AutoNormal.scales.bAR 0.079374395
AutoNormal.locs.sigma -2.205863
AutoNormal.scales.sigma 0.060526706

请注意，由于我们使用了较高的学习率，因此训练速度很快。有时模型和 guide 对学习率敏感，首先要尝试的是降低学习率并增加步数。这在包含深度神经网络的模型和 guide 中尤其重要。我们建议从较低的学习率开始，逐渐增加，避免过高的学习率，否则推断可能会发散或导致 NAN。

训练好 guide 后，我们可以从 Pyro 的 param 存储中获取优化后的 guide 参数值进行检查。下面打印的每个 (loc, scale) 对参数化了 guide 中对应于模型中不同未观测 pyro.sample 语句的单个 pyro.distributions.Normal 分布，这与我们之前手写的 custom_guide 类似。

[57]:

with pyro.plate("samples", 800, dim=-1):
    samples = auto_guide(is_cont_africa, ruggedness)

gamma_within_africa = samples["bR"] + samples["bAR"]
gamma_outside_africa = samples["bR"]

最后，让我们重新审视我们先前提出的问题：地形崎岖度和 GDP 之间的关系对于模型参数估计中的任何不确定性有多大的鲁棒性。为此，我们绘制了非洲内部和外部国家对数 GDP 与地形崎岖度斜率的分布。

[58]:

fig = plt.figure(figsize=(10, 6))
sns.histplot(gamma_within_africa.detach().cpu().numpy(), kde=True, stat="density", label="African nations")
sns.histplot(gamma_outside_africa.detach().cpu().numpy(), kde=True, stat="density", label="Non-African nations", color="orange")
fig.suptitle("Density of Slope : log(GDP) vs. Terrain Ruggedness");
plt.xlabel("Slope of regression line")
plt.legend()
plt.show()
_images/intro_long_59_0.png

我们使用从训练好的 guide 中抽取的样本来表示这两个分布。为了并行抽取多个样本，我们可以在 pyro.plate 语句中调用 guide，该语句重复并向量化 guide 中每个 pyro.sample 语句的采样操作，如介绍 pyro.plate 原语一节所述。

如下图所示，非洲国家的概率质量主要集中在正向区域，而其他国家则反之，这进一步证实了原始假设。然而，非非洲国家的后验不确定性（橙色直方图的宽度）似乎明显低于非洲国家（蓝色直方图的宽度），考虑到原始数据中看似相似的散布，这令人惊讶。我们将在下一节进一步探讨这一差异。

Pyro 中的模型评估

背景：使用后验预测检验进行贝叶斯模型评估

为了评估我们推断结果的可信度，我们将比较由模型诱导的可能新数据的后验预测分布与现有观测数据。通常情况下，计算这个分布是难以处理的，因为它依赖于已知真实的后验分布，但我们可以使用从变分推断获得的近似后验分布来轻松近似它

\[p_{\theta}(x' | {\bf x}) = \int \! d{\bf z}\; p_{\theta}(x' | {\bf z}) p_{\theta}({\bf z} | {\bf x}) \approx \int \! d{\bf z}\; p_{\theta}(x' | {\bf z}) q_{\phi}({\bf z} | {\bf x})\]

具体来说，要从后验预测分布中抽取一个近似样本，我们只需从近似后验分布中抽取一个样本 \({\hat {\bf z}} \sim q_{\phi}({\bf z})\)，然后根据模型中给定该样本的观测变量分布 \(x' \sim p_{\theta}(x | {\hat {\bf z}})\) 进行采样，就好像我们用近似后验分布替换了先验分布一样。

示例：Pyro 中的后验预测不确定性

[59]:

predictive = pyro.infer.Predictive(model, guide=auto_guide, num_samples=800)
svi_samples = predictive(is_cont_africa, ruggedness, log_gdp=None)
svi_gdp = svi_samples["obs"]

为了评估我们的线性回归模型示例，我们将使用 Predictive 工具类生成和可视化来自后验预测分布的一些样本，该类实现了上述从 \(p_{\theta}(x' | {\bf x})\) 中近似采样的步骤。

[60]:

predictions = pd.DataFrame({
    "cont_africa": is_cont_africa,
    "rugged": ruggedness,
    "y_mean": svi_gdp.mean(0).detach().cpu().numpy(),
    "y_perc_5": svi_gdp.kthvalue(int(len(svi_gdp) * 0.05), dim=0)[0].detach().cpu().numpy(),
    "y_perc_95": svi_gdp.kthvalue(int(len(svi_gdp) * 0.95), dim=0)[0].detach().cpu().numpy(),
    "true_gdp": log_gdp,
})
african_nations = predictions[predictions["cont_africa"] == 1].sort_values(by=["rugged"])
non_african_nations = predictions[predictions["cont_africa"] == 0].sort_values(by=["rugged"])

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6), sharey=True)
fig.suptitle("Posterior predictive distribution with 90% CI", fontsize=16)

ax[0].plot(non_african_nations["rugged"], non_african_nations["y_mean"])
ax[0].fill_between(non_african_nations["rugged"], non_african_nations["y_perc_5"], non_african_nations["y_perc_95"], alpha=0.5)
ax[0].plot(non_african_nations["rugged"], non_african_nations["true_gdp"], "o")
ax[0].set(xlabel="Terrain Ruggedness Index", ylabel="log GDP (2000)", title="Non African Nations")

ax[1].plot(african_nations["rugged"], african_nations["y_mean"])
ax[1].fill_between(african_nations["rugged"], african_nations["y_perc_5"], african_nations["y_perc_95"], alpha=0.5)
ax[1].plot(african_nations["rugged"], african_nations["true_gdp"], "o")
ax[1].set(xlabel="Terrain Ruggedness Index", ylabel="log GDP (2000)", title="African Nations");
_images/intro_long_64_0.png

我们从训练好的模型中生成了 800 个样本。在内部，这是通过首先从 guide 中生成潜变量的样本，然后向前运行模型，同时将未观测 pyro.sample 语句返回的值更改为从 guide 中采样的相应值来实现的。

下面的代码特定于本示例，仅用于绘制每个国家后验预测分布的 90% 可信区间（包含 90% 概率质量的区间）。

我们观察到，我们模型的结果和 90% CI 解释了我们在实践中观察到的大多数数据点，但仍然有相当多的非非洲国家被我们的近似后验认为是不太可能的。

示例：使用全秩 guide 重新审视贝叶斯回归
为了改善我们的结果，我们将尝试使用一个从所有参数的多变量正态分布中生成样本的 guide。这使我们能够通过全秩协方差矩阵 \(\Sigma \in \mathbb{R}^{5 \times 5}\) 捕捉潜变量之间的相关性；我们之前的 guide 忽略了这些相关性。也就是说，我们有

\[\alpha, \beta_a, \beta_r, \beta_{ar}, \sigma_u \sim q_{\phi = ({\bf \mu}, {\bf \Sigma})}(\alpha, \beta_a, \beta_r, \beta_{ar}, \sigma_u) = \rm{Normal}((\alpha, \beta_a, \beta_r, \beta_{ar}, \sigma_u) | {\bf \mu}, {\bf \Sigma})\]

\[\sigma = \rm{constrain}(\sigma_u)\]

[61]:

mvn_guide = pyro.infer.autoguide.AutoMultivariateNormal(model)

要手动编写这种形式的 guide，我们需要合并所有潜变量，以便从单个 pyro.distributions.MultivariateNormal 分布中对它们进行 pyro.sample，选择一个 \(\rm{constrain}()\) 的实现来固定 \(\sigma\) 的值为正，创建并初始化适当形状的参数 \(\mu, \Sigma\)，并约束变分参数 \(\Sigma\) 在整个优化过程中保持有效的协方差矩阵（即保持正定）。

[62]:

pyro.render_model(mvn_guide, model_args=(is_cont_africa, ruggedness, log_gdp), render_params=True)

[62]:
_images/intro_long_69_0.svg

这将相当繁琐，因此我们将使用另一个 autoguide 来替我们完成所有这些繁琐的工作，即 pyro.infer.autoguide.AutoMultivariateNormal

使用 pyro.render_model 表明，与我们的平均场 AutoNormal guide 不同，此 guide 明确捕捉了我们模型中所有潜变量之间的相关性。可视化图中的新节点 _AutoMultivariateNormal_latent 对应于上面的方程；其他对应于模型变量的节点仅索引到此张量值随机变量的单个元素中。

[63]:

%%time
pyro.clear_param_store()
mvn_guide = pyro.infer.autoguide.AutoMultivariateNormal(model)
svi = pyro.infer.SVI(model,
                     mvn_guide,
                     pyro.optim.Adam({"lr": 0.02}),
                     pyro.infer.Trace_ELBO())

losses = []
for step in range(1000 if not smoke_test else 2):
    loss = svi.step(is_cont_africa, ruggedness, log_gdp)
    losses.append(loss)
    if step % 100 == 0:
        logging.info("Elbo loss: {}".format(loss))

plt.figure(figsize=(5, 2))
plt.plot(losses)
plt.xlabel("SVI step")
plt.ylabel("ELBO loss")

with pyro.plate("samples", 800, dim=-1):
    mvn_samples = mvn_guide(is_cont_africa, ruggedness)

mvn_gamma_within_africa = mvn_samples["bR"] + mvn_samples["bAR"]
mvn_gamma_outside_africa = mvn_samples["bR"]

# Interface note: reuse guide samples for prediction by passing them to Predictive
# via the posterior_samples keyword argument instead of passing the guide as above
assert "obs" not in mvn_samples
mvn_predictive = pyro.infer.Predictive(model, posterior_samples=mvn_samples)
mvn_predictive_samples = mvn_predictive(is_cont_africa, ruggedness, log_gdp=None)

mvn_gdp = mvn_predictive_samples["obs"]

Elbo loss: 702.4906432628632
Elbo loss: 548.7575962543488
Elbo loss: 490.9642730951309
Elbo loss: 401.81392109394073
Elbo loss: 333.7779414653778
Elbo loss: 247.01823914051056
Elbo loss: 248.3894298672676
Elbo loss: 247.3512134552002
Elbo loss: 248.2095948457718
Elbo loss: 247.21006780862808

CPU times: user 1min 45s, sys: 21.9 ms, total: 1min 45s
Wall time: 7.03 s
_images/intro_long_71_2.png

我们的模型以及其余的推断和评估代码与之前基本相同：我们使用 pyro.optim.Adam 和 pyro.infer.Trace_ELBO 来拟合新 guide 的参数，然后从 guide 中采样并使用 Predictive 从后验预测分布中采样。

有一个值得注意的微小差异：我们通过 Predictive 的 posterior_samples 关键字参数将 guide 样本直接传递给 Predictive，而不是像上一节那样传递 guide。这避免了不必要的重复计算。

[64]:

svi_samples = {k: v.detach().cpu().numpy() for k, v in samples.items()}
svi_mvn_samples = {k: v.detach().cpu().numpy() for k, v in mvn_samples.items()}

fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))
fig.suptitle("Cross-sections of the Posterior Distribution", fontsize=16)
sns.kdeplot(x=svi_samples["bA"], y=svi_samples["bR"], ax=axs[0], bw_adjust=4 )
sns.kdeplot(x=svi_mvn_samples["bA"], y=svi_mvn_samples["bR"], ax=axs[0], shade=True, bw_adjust=4)
axs[0].set(xlabel="bA", ylabel="bR", xlim=(-2.8, -0.9), ylim=(-0.6, 0.2))

sns.kdeplot(x=svi_samples["bR"], y=svi_samples["bAR"], ax=axs[1],bw_adjust=4 )
sns.kdeplot(x=svi_mvn_samples["bR"], y=svi_mvn_samples["bAR"], ax=axs[1], shade=True, bw_adjust=4)
axs[1].set(xlabel="bR", ylabel="bAR", xlim=(-0.55, 0.2), ylim=(-0.15, 0.85))


for label, color in zip(["SVI (Diagonal Normal)", "SVI (Multivariate Normal)"], sns.color_palette()[:2]):
    plt.plot([], [],
                label=label, color=color)

fig.legend(loc='upper right')

[64]:

<matplotlib.legend.Legend at 0x7f8971b854c0>
_images/intro_long_74_1.png

现在让我们比较之前由 AutoDiagonalNormal guide 计算的后验分布与 AutoMultivariateNormal guide 计算的后验分布。我们将视觉上叠加后验分布的截面图（成对回归系数的联合分布）。

[65]:

fig, axs = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))
fig.suptitle("Density of Slope : log(GDP) vs. Terrain Ruggedness");

sns.histplot(gamma_within_africa.detach().cpu().numpy(), ax=axs[0], kde=True, stat="density", label="African nations")
sns.histplot(gamma_outside_africa.detach().cpu().numpy(), ax=axs[0], kde=True, stat="density", color="orange", label="Non-African nations")
axs[0].set(title="Mean field", xlabel="Slope of regression line", xlim=(-0.6, 0.6), ylim=(0, 11))

sns.histplot(mvn_gamma_within_africa.detach().cpu().numpy(), ax=axs[1], kde=True, stat="density", label="African nations")
sns.histplot(mvn_gamma_outside_africa.detach().cpu().numpy(), ax=axs[1], kde=True, stat="density", color="orange", label="Non-African nations")
axs[1].set(title="Full rank", xlabel="Slope of regression line", xlim=(-0.6, 0.6), ylim=(0, 11))

handles, labels = axs[1].get_legend_handles_labels()
fig.legend(handles, labels, loc='upper right');
_images/intro_long_76_0.png

请注意，多变量正态近似比平均场近似更分散，并且能够模拟后验分布中系数之间的相关性。

[66]:

mvn_predictions = pd.DataFrame({
    "cont_africa": is_cont_africa,
    "rugged": ruggedness,
    "y_mean": mvn_gdp.mean(dim=0).detach().cpu().numpy(),
    "y_perc_5": mvn_gdp.kthvalue(int(len(mvn_gdp) * 0.05), dim=0)[0].detach().cpu().numpy(),
    "y_perc_95": mvn_gdp.kthvalue(int(len(mvn_gdp) * 0.95), dim=0)[0].detach().cpu().numpy(),
    "true_gdp": log_gdp,
})
mvn_non_african_nations = mvn_predictions[mvn_predictions["cont_africa"] == 0].sort_values(by=["rugged"])

fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6), sharey=True)
fig.suptitle("Posterior predictive distribution with 90% CI", fontsize=16)

ax[0].plot(non_african_nations["rugged"], non_african_nations["y_mean"])
ax[0].fill_between(non_african_nations["rugged"], non_african_nations["y_perc_5"], non_african_nations["y_perc_95"], alpha=0.5)
ax[0].plot(non_african_nations["rugged"], non_african_nations["true_gdp"], "o")
ax[0].set(xlabel="Terrain Ruggedness Index", ylabel="log GDP (2000)", title="Non African Nations: Mean-field")

ax[1].plot(mvn_non_african_nations["rugged"], mvn_non_african_nations["y_mean"])
ax[1].fill_between(mvn_non_african_nations["rugged"], mvn_non_african_nations["y_perc_5"], mvn_non_african_nations["y_perc_95"], alpha=0.5)
ax[1].plot(mvn_non_african_nations["rugged"], mvn_non_african_nations["true_gdp"], "o")
ax[1].set(xlabel="Terrain Ruggedness Index", ylabel="log GDP (2000)", title="Non-African Nations: Full rank");
_images/intro_long_78_0.png

我们可以通过重复我们之前对非洲内部和外部国家崎岖度-GDP 系数分布的可视化来看到这一点的影响。现在，这两个系数的后验不确定性大致相同，这与目测数据所显示的情况一致。

我们在两种近似下可视化了非非洲国家后验预测分布的 90% 可信区间，验证了我们对观测数据的覆盖范围有所改善

后续步骤

如果你已经看到这里，你就可以开始使用 Pyro 了！按照首页的说明安装 Pyro，并查看我们其他的示例和教程，尤其是实用 Pyro 和 PyTorch 系列教程，其中包括本教程中同一贝叶斯回归分析的一个版本，该版本使用更PyTorch 原生建模 API 编写。