tensorflow 概率的体系结构 MultivariateNormalFullCovariance

Question

努力学习tfp脑放屁。我想在给定一些解释性输入的情况下对数据集的均值和协方差进行建模。我知道我的架构是错误的，但我不确定如何。请看下面的代码：

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

# Getting things set up
tfd = tfp.distributions
x_train = np.random.normal(0,1,(1000, 20))
y_train =  np.random.normal(0,1,(1000, 10))

# Define the prior and posterior distributions
def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
  n = kernel_size + bias_size
  c = np.log(np.expm1(1.))
  return tf.keras.Sequential([
      tfp.layers.VariableLayer(2 * n, dtype=dtype),
      tfp.layers.DistributionLambda(lambda t: tfd.MultivariateNormalFullCovariance(
          loc=t[..., :n],
          covariance_matrix=tfp.math.fill_triangular(t[..., n:], upper=True) ** 2)),
  ])
def prior_trainable(kernel_size, bias_size=0, dtype=None):
  n = kernel_size + bias_size
  return tf.keras.Sequential([
      tfp.layers.VariableLayer(n + n*(n+1)//2, dtype=dtype),
      tfp.layers.DistributionLambda(lambda t: tfd.MultivariateNormalFullCovariance(
          loc=t[..., :n],
          covariance_matrix=tfp.math.fill_triangular(t[..., n:], upper=True) ** 2)),
  ])

# Define the model architecture
model = tf.keras.Sequential([
  tfp.layers.DenseVariational(2 + 2, posterior_mean_field, prior_trainable, kl_weight=1/x_train.shape[0]),
  tfp.layers.DistributionLambda(
      lambda t: tfd.MultivariateNormalFullCovariance(
          loc=t[..., :2],
          covariance_matrix=tfp.math.fill_triangular(t[..., 2:], upper=True) ** 2)),
])
model
# Define the negative log-likelihood loss function
def negloglik(y_true, y_pred):
    return -y_pred.log_prob(y_true)

# Compile the model with the negative log-likelihood loss function
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01),
              loss=negloglik)

# Train the model on some data
model.fit(x_train, y_train, epochs=10)

# Make predictions and compute the mean and covariance matrix
x_test = tf.random.normal((10, 10))
predicted_distribution = model.predict(x_test)
mean = predicted_distribution.mean()
cov = predicted_distribution.covariance()

print('Mean:', mean.numpy())
print('Covariance matrix:', cov.numpy())

运行该代码会产生以下错误： ValueError：调用层“distribution_lambda”（类型 DistributionLambda）时遇到异常。

Input right-most shape (84) does not correspond to a triangular matrix.

Call arguments received by layer 'distribution_lambda' (type DistributionLambda):
  • inputs=tf.Tensor(shape=(168,), dtype=float32)
  • args=<class 'inspect._empty'>
  • kwargs={'training': 'True'}

谢谢！