python - 张量流概率中流归一化的变分推理

Question

上次我读到了一些关于使用规范化流来改进变分推理 fe Link1 Link2的内容。

Tensorflow 概率已经在 bijector 子模块中提供了 RealNVP 和 MaskedAutoregressiveFlow，并且在 layers 子模块中提供了 AutoregressiveTransform 层。因此，我认为构建一个经过变分推理训练的贝叶斯神经网络和使用 Tensorflow Probability 的归一化流给出的后验将是简单而直接的。

从其中一个教程（链接）开始，我能够构建一个带有 mean_field_posterior 的 BNN。

然后事情开始变得复杂起来。我编写了以下函数，改编自此示例 ( Link )，以生成使用掩码自回归流转换正态分布的后验。

def posterior_vi_maf(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    return tf.keras.Sequential(
        [
            tfk.layers.InputLayer(input_shape=(0,), dtype=tf.float32),
            tfpl.DistributionLambda(
                lambda t: tfd.MultivariateNormalDiag(
                    loc=tf.zeros(tf.concat([tf.shape(t)[:-1], [4]], axis=0)),
                    scale_diag=tf.ones(4),
                ),
            ),
            tfp.layers.AutoregressiveTransform(
                tfb.AutoregressiveNetwork(
                    params=2, hidden_units=[10, 10], activation="relu"
                )
            ),
        ]
    )

比较terior_vi_maf 和terior_mean_field 的形状和输出，从技术角度来看，似乎一切都应该起作用。

p1 = posterior_vi_maf(16, 4, dtype=tf.float32)
p2 = posterior_mean_field(16, 4, dtype=tf.float32)
assert p1(x).shape == p2(x).shape
assert isinstance(p1(x), tfd.Distribution)
assert isinstance(p2(x), tfd.Distribution)

不幸的是，运行训练脚本（见底部）会引发以下错误消息：

ValueError: Shape must be rank 1 but is rank 2 for '{{node dense_variational/BiasAdd}} = BiasAdd[T=DT_FLOAT, data_format="NHWC"](dense_variational/MatMul, dense_variational/split:1)' with input shapes: [?,?,16], [?,16].

有什么建议为什么会发生这种情况和/或我该如何解决这个问题？

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfk = tf.keras
tfl = tf.keras.layers
tfpl = tfp.layers
tfb = tfp.bijectors

np.random.seed(2)
tf.random.set_seed(2)
N = 100

x = tfd.Normal(loc=0, scale=1).sample(N)
y = tfd.Normal(loc=x * 0.5, scale=0.3).sample()


def negloglik(y_true, y_pred):
    nll = -tf.reduce_mean(y_pred.log_prob(y_true))
    return nll


def prior(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    return lambda t: tfd.Independent(
        tfd.Normal(loc=tf.zeros(n, dtype=dtype), scale=1),
        reinterpreted_batch_ndims=1,
    )


def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    c = np.log(np.expm1(1.0))
    return tf.keras.Sequential(
        [
            tfp.layers.VariableLayer(2 * n, dtype=dtype),
            tfp.layers.DistributionLambda(
                lambda t: tfd.Independent(
                    tfd.Normal(
                        loc=t[..., :n],
                        scale=1e-5 + 0.01 * tf.nn.softplus(c + t[..., n:]),
                    ),
                    reinterpreted_batch_ndims=1,
                )
            ),
        ]
    )


model = tf.keras.Sequential(name="small_vi_nn")
model.add(tfl.Input(1))
model.add(
    tfpl.DenseVariational(
        units=16,
        make_posterior_fn=posterior_vi_maf,
        make_prior_fn=prior,
        kl_weight=1 / N,
        kl_use_exact=False,
        activation="relu",
    )
)
model.add(
    tfpl.DenseVariational(
        units=2,
        make_posterior_fn=posterior_mean_field,
        make_prior_fn=prior,
        kl_weight=1 / N,
        kl_use_exact=False,
    )
)
model.add(
    tfpl.DistributionLambda(
        make_distribution_fn=lambda t: tfd.Normal(
            loc=t[:, 0],
            scale=1e-3 + tf.math.softplus(0.05 * t[:, 0]),
        )
    )
)

optimizer = tf.keras.optimizers.Adam(learning_rate=0.01, amsgrad=True)


model.compile(optimizer=optimizer, loss=negloglik)
model.fit(
    x,
    y,
    epochs=2,
    shuffle=True,
    verbose=True,
)

标签： pythonkerastensorflow-probabilitybayesian-deep-learningdensevariational