Tensorflow Probability 中标准化流的变分推理

问题描述

在上一次我读过一些关于使用规范化流来改进变分推理 f.e. Link1 Link2。

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

从其中一个教程 (Link) 开始，我能够构建一个带有 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"
                )
            ),]
    )

比较后验_vi_maf 和后验_平均_场的形状和输出，从技术角度来看，似乎一切都应该可行。

p1 = posterior_vi_maf(16,4,dtype=tf.float32)
p2 = posterior_mean_field(16,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,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,dtype=None):
    n = kernel_size + bias_size
    c = np.log(np.expm1(1.0))
    return tf.keras.Sequential(
        [
            tfp.layers.VariableLayer(2 * n,tfp.layers.distributionLambda(
                lambda t: tfd.Independent(
                    tfd.normal(
                        loc=t[...,:n],scale=1e-5 + 0.01 * tf.nn.softplus(c + t[...,n:]),)
            ),]
    )


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,)
)
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,)

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）

bayesian-deep-learning densevariational keras keras python tensorflow-probability