在张量流概率中构建基于 table 的离散 CPD?

Constructing discrete table-based CPDs in tensorflow-probablity?

我正在尝试构建具有多个离散随机变量和条件概率的贝叶斯网络的最简单示例(Koller 书中的 "Student Network",参见 1

虽然有点笨拙,但我设法使用pymc3 构建了这个网络。特别是,在 pymc3 中创建 CPD 并不是那么简单,请参见下面的代码片段:

import pymc3 as pm

...

with pm.Model() as basic_model:
    # parameters for categorical are indexed as [0, 1, 2, ...]
    difficulty = pm.Categorical(name='difficulty', p=[0.6, 0.4])

    intelligence = pm.Categorical(name='intelligence', p=[0.7, 0.3])

    grade = pm.Categorical(name='grade',
        p=pm.math.switch(
            theano.tensor.eq(intelligence, 0),
                pm.math.switch(
                    theano.tensor.eq(difficulty, 0),
                        [0.3, 0.4, 0.3],  # I=0, D=0
                        [0.05, 0.25, 0.7]   # I=0, D=1
                    ),
                    pm.math.switch(
                        theano.tensor.eq(difficulty, 0),
                            [0.9, 0.08, 0.02],  # I=1, D=0
                            [0.5, 0.3, 0.2]  # I=1, D=1
                    )
            )
        )

    letter = pm.Categorical(name='letter', p=pm.math.switch(
    ...

但我不知道如何使用 tensoflow-probability 构建这个网络(版本:tfp-nightly==0.7.0.dev20190517tf-nightly-2.0-preview==2.0.0.dev20190517

对于无条件二元变量,可以采用分类分布,如

from tensorflow_probability import distributions as tfd
from tensorflow_probability import edward2 as ed

difficulty = ed.RandomVariable(
                 tfd.Categorical(
                     probs=[0.6, 0.4],
                     name='difficulty'
                 )
             )

但是如何构建 CPD?

tensorflow-probability 中几乎没有 classes/methods 可能相关(在 tensorflow_probability/python/distributions/deterministic.py 或已弃用的 ConditionalDistribution 中),但文档相当稀疏(需要深入了解 tfp ).

--- 更新问题 ---

Chris 的回答是一个很好的起点。然而,即使对于一个非常简单的双变量模型,事情仍然有点不清楚。

效果很好:

jdn = tfd.JointDistributionNamed(dict(
    dist_x=tfd.Categorical([0.2, 0.8], validate_args=True),
    dist_y=lambda dist_x: tfd.Bernoulli(probs=tf.gather([0.1, 0.9], indices=dist_x), validate_args=True)
))
print(jdn.sample(10))

但是这个失败了

jdn = tfd.JointDistributionNamed(dict(
    dist_x=tfd.Categorical([0.2, 0.8], validate_args=True),
    dist_y=lambda dist_x: tfd.Categorical(probs=tf.gather_nd([[0.1, 0.9], [0.5, 0.5]], indices=[dist_x]))
))
print(jdn.sample(10))

(出于学习目的,我试图在第二个示例中明确地建模分类)

-- 更新:已解决 ---

显然,最后一个示例错误地使用了 tf.gather_nd 而不是 tf.gather,因为我们只想 select 基于 dist_x outome 的第一行或第二行。此代码现在有效:

jdn = tfd.JointDistributionNamed(dict(
    dist_x=tfd.Categorical([0.2, 0.8], validate_args=True),
    dist_y=lambda dist_x: tfd.Categorical(probs=tf.gather([[0.1, 0.9], [0.5, 0.5]], indices=[dist_x]))
))
print(jdn.sample(10))

这件事的棘手之处,大概也是它比 PyMC 中预期的更微妙的原因,是——与矢量化编程中的几乎所有事情一样——处理形状。

在 TF/TFP 中,(IMO)解决此问题的最佳方法是使用新的 TFP JointDistribution{Sequential,Named,Coroutine} 类 之一。这些让您自然地表示分层 PGM 模型,然后从中采样,评估对数概率等。

我为完整的学生网络创建了一个演示所有 3 种方法的 colab 笔记本:https://colab.research.google.com/drive/1D2VZ3OE6tp5pHTsnOAf_7nZZZ74GTeex

请注意 tf.gather 和 tf.gather_nd 的关键用途是管理各种二进制和分类切换的矢量化。

看看吧,如果有任何问题,请告诉我!