用于回归任务的贝叶斯 CNN

如何解决用于回归任务的贝叶斯 CNN

我有一个标准的 CNN 模型来解决图片数据集中的回归任务。该模型是使用 Tensorflow 实现的，并且在我的数据集上运行良好：

def create_cnn_model() -> Model:
    cnn = models.Sequential()
    cnn.add(layers.Conv2D(32,3,activation="relu",input_shape=(450,800,3)))
    cnn.add(layers.MaxPool2D((2,2)))
    cnn.add(layers.Conv2D(64,activation="relu"))
    cnn.add(layers.MaxPooling2D(2,2))
    cnn.add(layers.Conv2D(64,activation="relu"))

    cnn.add(layers.Flatten())

    cnn.add(layers.Dense(128,activation="relu"))
    cnn.add(layers.Dense(64,activation="relu"))
    cnn.add(layers.Dense(1,activation="linear"))

    cnn.compile(
        optimizer="adam",loss=tf.keras.losses.MeanSquaredError()
    )

    return cnn

现在我想将 CNN 更改为贝叶斯 CNN 模型，以便我可以计算预测的置信度以及区间（最小-最大预测）。为此，我尝试使用库 tensorflow-probability，如下所示：

def create_bcnn_model(train_size) -> Model:
    bcnn = models.Sequential()
    bcnn.add(layers.Conv2D(32,3)))
    bcnn.add(layers.MaxPool2D((2,2)))
    bcnn.add(layers.Conv2D(64,activation="relu"))
    bcnn.add(layers.MaxPooling2D(2,2))
    bcnn.add(layers.Conv2D(64,activation="relu"))

    bcnn.add(layers.Flatten())

    hidden_units = [128,64]
    for units in hidden_units:
        bcnn.add(tfp.layers.DenseVariational(
            units=units,make_prior_fn=prior,make_posterior_fn=posterior,kl_weight=1 / train_size,))

    bcnn.add(layers.Dense(1,activation="linear"))

    bcnn.compile(
        optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.0005),loss=tf.keras.losses.MeanSquaredError(),metrics=[tf.keras.metrics.RootMeanSquaredError()],)

    return bcnn

我基本上用 DenseVariational 层替换了标准的 Dense 层。我主要是按照 keras.io 中的教程进行操作。但我无法让它工作并收到以下错误消息：

ValueError: Input right-most shape (566353984) does not correspond to a triangular matrix.

不允许将标准 tensorflow 层与 tfp 层混合使用吗？有没有贝叶斯 CNN 用于回归任务的好例子？我能找到的所有例子都是关于分类的。我将不胜感激任何方向的任何提示或技巧:)

编辑：我的先验和后验与 keras 示例中的相同：

def prior(kernel_size,bias_size,dtype=None):
    n = kernel_size + bias_size
    prior_model = tf.keras.Sequential(
        [
            tfp.layers.distributionLambda(
                lambda t: tfp.distributions.MultivariatenormalDiag(
                    loc=tf.zeros(n),scale_diag=tf.ones(n)
                )
            )
        ]
    )
    return prior_model


def posterior(kernel_size,dtype=None):
    n = kernel_size + bias_size
    posterior_model = tf.keras.Sequential(
        [
            tfp.layers.VariableLayer(
                tfp.layers.MultivariatenormalTriL.params_size(n),dtype=dtype
            ),tfp.layers.MultivariatenormalTriL(n),]
    )
    return posterior_model