Edward2
Edward2是一个简单的概率编程语言。它在深度学习生态系统中提供核心实用工具,使得用户可以将模型编写为概率程序,并操作模型的计算以实现灵活的训练和推断。它的组织结构如下:
edward2/
: 库代码。examples/
: 示例。experimental/
: 活跃的研究项目。
你是否正在从Edward升级?请查看指南
Upgrading_from_Edward_to_Edward2.md
。
核心实用工具相当底层:如果你想要一个用于不确定性建模的高级模块,请查看
贝叶斯层的指南。
如果你想要基于研究就绪的代码进行开发,我们推荐使用
Uncertainty Baselines。
安装
我们推荐使用最新的开发版本。要安装,请运行:
pip install "edward2 @ git+https://github.com/google/edward2.git"
你也可以使用以下命令安装最新的稳定版本。但请注意,我们很少更新稳定版本(这是一个由兼职人员维护的热情项目,定期安排发布会占用大量时间)。
pip install edward2
Edward2支持三个后端:TensorFlow(默认)、JAX和NumPy(参见下文激活)。安装edward2
并不会自动安装任何后端。要获取这些依赖项,请使用例如pip install edward2[tensorflow]"
,将tensorflow
替换为相应的后端。有时Edward2会使用TensorFlow的最新变更,这种情况下你需要TensorFlow的每日构建包:使用pip install edward2[tf-nightly]
。
1. 模型作为概率程序
随机变量
在Edward2中,我们使用
RandomVariables
来指定概率模型的结构。
随机变量rv
携带一个概率分布(rv.distribution
),
这是一个TensorFlow Distribution实例,用于管理随机变量的方法,
如log_prob
和sample
。
随机变量的形式类似于TensorFlow Distributions。
import edward2 as ed
normal_rv = ed.Normal(loc=0., scale=1.)
## <ed.RandomVariable 'Normal/' shape=() dtype=float32 numpy=0.0024812892>
normal_rv.distribution.log_prob(1.231)
## <tf.Tensor: id=11, shape=(), dtype=float32, numpy=-1.6766189>
dirichlet_rv = ed.Dirichlet(concentration=tf.ones([2, 3]))
## <ed.RandomVariable 'Dirichlet/' shape=(2, 3) dtype=float32 numpy=
array([[0.15864784, 0.01217205, 0.82918006],
[0.23385087, 0.69622266, 0.06992647]], dtype=float32)>
默认情况下,实例化随机变量rv
会创建一个采样操作来形成
张量rv.value ~ rv.distribution.sample()
。默认的样本数量
(可通过rv
的sample_shape
参数控制)为1,如果提供了可选的value
参数,
则不会创建采样操作。随机变量可以与TensorFlow操作互操作:TF操作在样本上进行操作。
x = ed.Normal(loc=tf.zeros(2), scale=tf.ones(2))
y = 5.
x + y, x / y
## (<tf.Tensor: id=109, shape=(2,), dtype=float32, numpy=array([3.9076924, 4.588356 ], dtype=float32)>,
## <tf.Tensor: id=111, shape=(2,), dtype=float32, numpy=array([-0.21846154, -0.08232877], dtype=float32)>)
tf.tanh(x * y)
## <tf.Tensor: id=114, shape=(2,), dtype=float32, numpy=array([-0.99996394, -0.9679181 ], dtype=float32)>
x[1] # 第2个正态随机变量
## <ed.RandomVariable 'Normal/' shape=() dtype=float32 numpy=-0.41164386>
概率模型
Edward2中的概率模型表示为实例化一个或多个RandomVariables
的Python函数。
通常,该函数("程序")执行生成过程并返回样本。函数的输入可以被视为模型条件的值。
下面我们编写贝叶斯逻辑回归,其中给定特征、系数和截距生成二元结果。 系数和截距有先验分布。执行该函数会添加操作,从先验中采样系数和截距, 并使用这些样本来计算结果。
def logistic_regression(features):
"""贝叶斯逻辑回归 p(y | x) = int p(y | x, w, b) p(w, b) dwdb。"""
coeffs = ed.Normal(loc=tf.zeros(features.shape[1]), scale=1., name="coeffs")
intercept = ed.Normal(loc=0., scale=1., name="intercept")
outcomes = ed.Bernoulli(
logits=tf.tensordot(features, coeffs, [[1], [0]]) + intercept,
name="outcomes")
return outcomes
num_features = 10
features = tf.random.normal([100, num_features])
outcomes = logistic_regression(features)
# <ed.RandomVariable 'outcomes/' shape=(100,) dtype=int32 numpy=
# array([1, 0, ... 0, 1], dtype=int32)>
Edward2程序还可以表示超出直接建模数据的分布。例如,下面我们编写一个可学习的分布, 目的是将其近似到逻辑回归的后验分布。
def logistic_regression_posterior(coeffs_loc, coeffs_scale,
intercept_loc, intercept_scale):
"""贝叶斯逻辑回归的后验分布 p(w, b | {x, y})。"""
coeffs = ed.MultivariateNormalTriL(
loc=coeffs_loc,
scale_tril=tfp.trainable_distributions.tril_with_diag_softplus_and_shift(
coeffs_scale),
name="coeffs_posterior")
intercept = ed.Normal(
loc=intercept_loc,
scale=tf.nn.softplus(intercept_scale) + 1e-5,
name="intercept_posterior")
return coeffs, intercept
coeffs_loc = tf.Variable(tf.random.normal([num_features]))
coeffs_scale = tf.Variable(tf.random.normal(
[num_features*(num_features+1) // 2]))
intercept_loc = tf.Variable(tf.random.normal([]))
intercept_scale = tf.Variable(tf.random.normal([]))
posterior_coeffs, posterior_intercept = logistic_regression_posterior(
coeffs_loc, coeffs_scale, intercept_loc, intercept_scale)
2. 操作模型计算
跟踪
训练和测试概率模型通常需要的不仅仅是生成过程的样本。为了实现灵活的训练和测试, 我们使用跟踪来操作模型的计算。
跟踪器是一个作用于另一个函数f
及其参数*args
、**kwargs
的函数。
它在返回输出(通常是f(*args, **kwargs)
:应用函数本身的结果)之前执行各种计算。
ed.trace
上下文管理器将跟踪器推入堆栈,任何可跟踪的函数都会被堆栈拦截。
所有随机变量构造函数都是可跟踪的。
下面我们跟踪逻辑回归模型的生成过程。特别是,我们使用学习到的后验均值而不是先验来进行预测。
def set_prior_to_posterior_mean(f, *args, **kwargs):
"""形成后验预测,将每个先验设置为其后验均值。"""
name = kwargs.get("name")
if name == "coeffs":
return posterior_coeffs.distribution.mean()
elif name == "intercept":
return posterior_intercept.distribution.mean()
return f(*args, **kwargs)
with ed.trace(set_prior_to_posterior_mean):
predictions = logistic_regression(features)
training_accuracy = (
tf.reduce_sum(tf.cast(tf.equal(predictions, outcomes), tf.float32)) /
tf.cast(outcomes.shape[0], tf.float32))
程序转换
使用跟踪,还可以应用程序转换,将模型从一种表示映射到另一种表示。
这为根据下游用例访问不同的模型属性提供了便利。
例如,马尔可夫链蒙特卡洛算法通常需要模型的对数联合概率函数作为输入。下面我们以贝叶斯逻辑回归程序为例,该程序指定了一个生成过程,并应用内置的ed.make_log_joint
转换来获取其对数联合概率函数。对数联合函数将生成程序的原始输入以及程序中的随机变量作为输入。它返回一个标量张量,对所有随机变量的对数概率进行求和。
在我们的例子中,features
和outcomes
是固定的,我们想使用哈密顿蒙特卡洛方法从coeffs
和intercept
的后验分布中抽取样本。为此,我们创建了target_log_prob_fn
,它只将coeffs
和intercept
作为参数,并将输入features
和输出随机变量outcomes
固定为其已知值。
import no_u_turn_sampler # 本地文件导入
# 设置训练数据
features = tf.random.normal([100, 55])
outcomes = tf.random.uniform([100], minval=0, maxval=2, dtype=tf.int32)
# 将目标对数概率函数传递给MCMC转换核
log_joint = ed.make_log_joint_fn(logistic_regression)
def target_log_prob_fn(coeffs, intercept):
"""目标对数概率作为状态的函数。"""
return log_joint(features,
coeffs=coeffs,
intercept=intercept,
outcomes=outcomes)
coeffs_samples = []
intercept_samples = []
coeffs = tf.random.normal([55])
intercept = tf.random.normal([])
target_log_prob = None
grads_target_log_prob = None
for _ in range(1000):
[
[coeffs, intercepts],
target_log_prob,
grads_target_log_prob,
] = no_u_turn_sampler.kernel(
target_log_prob_fn=target_log_prob_fn,
current_state=[coeffs, intercept],
step_size=[0.1, 0.1],
current_target_log_prob=target_log_prob,
current_grads_target_log_prob=grads_target_log_prob)
coeffs_samples.append(coeffs)
intercept_samples.append(coeffs)
返回的coeffs_samples
和intercept_samples
分别包含coeffs
和intercept
的1,000个后验样本。例如,它们可以用于评估模型在新数据上的后验预测。
使用JAX或NumPy后端
使用替代后端非常简单,如下所示:
import edward2.numpy as ed # NumPy后端
import edward2.jax as ed # 或者,JAX后端
在NumPy后端中,Edward2封装了SciPy分布。例如,这里是线性回归。
def linear_regression(features, prior_precision):
beta = ed.norm.rvs(loc=0.,
scale=1. / np.sqrt(prior_precision),
size=features.shape[1])
y = ed.norm.rvs(loc=np.dot(features, beta), scale=1., size=1)
return y
参考文献
通常,我们建议引用以下文章。
Tran, D., Hoffman, M. D., Moore, D., Suter, C., Vasudevan S., Radul A., Johnson M., and Saurous R. A. (2018). 简单、分布式和加速的概率编程。 发表于 神经信息处理系统。
@inproceedings{tran2018simple,
author = {Dustin Tran and Matthew D. Hoffman and Dave Moore and Christopher Suter and Srinivas Vasudevan and Alexey Radul and Matthew Johnson and Rif A. Saurous},
title = {Simple, Distributed, and Accelerated Probabilistic Programming},
booktitle = {Neural Information Processing Systems},
year = {2018},
}
如果您想特别引用层模块,请使用以下文章。
Tran, D., Dusenberry M. W., van der Wilk M., Hafner D. (2019). 贝叶斯层:神经网络不确定性的模块。 发表于 神经信息处理系统。
@inproceedings{tran2019bayesian,
author = {Dustin Tran and Michael W. Dusenberry and Danijar Hafner and Mark van der Wilk},
title={Bayesian {L}ayers: A module for neural network uncertainty},
booktitle = {Neural Information Processing Systems},
year={2019}
}