使用tfp.layers.DenseVariational时遇tuple无rank属性的AttributeError

问题场景 #

基于Keras官方示例构建贝叶斯神经网络，添加tfp.layers.DenseVariational贝叶斯层时触发以下错误：

File "/Users/S/Documents/B/Prediction/test1.py", line 148, in <module>
    bnn_model_small = create_bnn_model(train_sample_size)
  File "/Users/S/Documents/B/Prediction/test1.py", line 131, in create_bnn_model
    features = tfp.layers.DenseVariational(
  File "/Users/S/.local/share/virtualenvs/B-hz56sUDM/lib/python3.9/site-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/S/.local/share/virtualenvs/B-hz56sUDM/lib/python3.9/site-packages/tf_keras/src/engine/input_spec.py", line 251, in assert_input_compatibility
    ndim = x.shape.rank
AttributeError: 'tuple' object has no attribute 'rank'

曾尝试使用from tf_agents.environments import tf_py_environment; environment = tf_py_environment.TFPyEnvironment(environment)解决，但因要求numpy<1.20引发大量依赖冲突。

完整代码如下：

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_datasets as tfds
import tensorflow_probability as tfp

dataset_size = 4898
batch_size = 256
train_size = int(dataset_size * 0.85)

# Create training and evaluation datasets
def get_train_and_test_splits(train_size, batch_size=1):
    dataset = (
        tfds.load(name="wine_quality", as_supervised=True, split="train")
        .map(lambda x, y: (x, tf.cast(y, tf.float32)))
        .prefetch(buffer_size=dataset_size)
        .cache()
    )
    train_dataset = (
        dataset.take(train_size).shuffle(buffer_size=train_size).batch(batch_size)
    )
    test_dataset = dataset.skip(train_size).batch(batch_size)
    return train_dataset, test_dataset

train_dataset, test_dataset = get_train_and_test_splits(train_size, batch_size)

hidden_units = [8, 8]
learning_rate = 0.001
num_epochs = 100
mse_loss = keras.losses.MeanSquaredError()

def run_experiment(model, loss, train_dataset, test_dataset):
    model.compile(
        optimizer=keras.optimizers.RMSprop(learning_rate=learning_rate),
        loss=loss,
        metrics=[keras.metrics.RootMeanSquaredError()],
    )
    print("Start training the model...")
    model.fit(train_dataset, epochs=num_epochs, validation_data=test_dataset)
    print("Model training finished.")
    _, rmse = model.evaluate(train_dataset, verbose=0)
    print(f"Train RMSE: {round(rmse, 3)}")
    print("Evaluating model performance...")
    _, rmse = model.evaluate(test_dataset, verbose=0)
    print(f"Test RMSE: {round(rmse, 3)}")

FEATURE_NAMES = [
    "fixed acidity", "volatile acidity", "citric acid",
    "residual sugar", "chlorides", "free sulfur dioxide",
    "total sulfur dioxide", "density", "pH", "sulphates", "alcohol",
]

def create_model_inputs():
    inputs = {}
    for feature_name in FEATURE_NAMES:
        inputs[feature_name] = layers.Input(
            name=feature_name, shape=(1,), dtype=tf.float32
        )
    return inputs

def prior(kernel_size, bias_size, dtype=None):
    n = kernel_size + bias_size
    prior_model = 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, bias_size, dtype=None):
    n = kernel_size + bias_size
    posterior_model = keras.Sequential(
        [
            tfp.layers.VariableLayer(
                tfp.layers.MultivariateNormalTriL.params_size(n), dtype=dtype
            ),
            tfp.layers.MultivariateNormalTriL(n),
        ]
    )
    return posterior_model

def create_bnn_model(train_size):
    inputs = create_model_inputs()
    features = keras.layers.concatenate(list(inputs.values()))
    features = layers.BatchNormalization()(features)

    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1 / train_size,
            activation="sigmoid",
        )(features)

    outputs = layers.Dense(units=1)(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

num_epochs = 500
train_sample_size = int(train_size * 0.3)
small_train_dataset = train_dataset.unbatch().take(train_sample_size).batch(batch_size)

bnn_model_small = create_bnn_model(train_sample_size)
run_experiment(bnn_model_small, mse_loss, small_train_dataset, test_dataset)

sample = 10
examples, targets = list(test_dataset.unbatch().shuffle(batch_size * 10).batch(sample))[0]

def compute_predictions(model, iterations=100):
    predicted = []
    for _ in range(iterations):
        predicted.append(model(examples).numpy())
    predicted = np.concatenate(predicted, axis=1)

    prediction_mean = np.mean(predicted, axis=1).tolist()
    prediction_min = np.min(predicted, axis=1).tolist()
    prediction_max = np.max(predicted, axis=1).tolist()
    prediction_range = (np.max(predicted, axis=1) - np.min(predicted, axis=1)).tolist()

    for idx in range(sample):
        print(
            f"Predictions mean: {round(prediction_mean[idx], 2)}, "
            f"min: {round(prediction_min[idx], 2)}, "
            f"max: {round(prediction_max[idx], 2)}, "
            f"range: {round(prediction_range[idx], 2)} - "
            f"Actual: {targets[idx]}"
        )

compute_predictions(bnn_model_small)

num_epochs = 500
bnn_model_full = create_bnn_model(train_size)
run_experiment(bnn_model_full, mse_loss, train_dataset, test_dataset)

compute_predictions(bnn_model_full)

错误原因 #

该错误源于TFP的DenseVariational层与KerasBatchNormalization层的输出不兼容：在较新版本的Keras中，BatchNormalization层如果未明确指定training参数，会返回包含训练状态的tuple（而非单一张量），而DenseVariational层期望接收纯张量输入，导致无法读取shape.rank属性。

解决方案 #

有两种可行的修复方式：

方案1：明确指定BatchNormalization的training参数 #

修改create_bnn_model函数中的BatchNormalization调用，添加training=False参数，确保输出为单一张量：

def create_bnn_model(train_size):
    inputs = create_model_inputs()
    features = keras.layers.concatenate(list(inputs.values()))
    # 明确指定training=False，避免返回tuple
    features = layers.BatchNormalization()(features, training=False)

    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1 / train_size,
            activation="sigmoid",
        )(features)

    outputs = layers.Dense(units=1)(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

方案2：改用LayerNormalization替代BatchNormalization #

LayerNormalization不需要维护全局统计量，输出始终是单一张量，更适合与贝叶斯层配合：

def create_bnn_model(train_size):
    inputs = create_model_inputs()
    features = keras.layers.concatenate(list(inputs.values()))
    # 替换为LayerNormalization
    features = layers.LayerNormalization()(features)

    for units in hidden_units:
        features = tfp.layers.DenseVariational(
            units=units,
            make_prior_fn=prior,
            make_posterior_fn=posterior,
            kl_weight=1 / train_size,
            activation="sigmoid",
        )(features)

    outputs = layers.Dense(units=1)(features)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

额外建议 #

确保TFP与TensorFlow版本匹配（例如TFP 0.20.x对应TensorFlow 2.10+），避免因API版本差异引发兼容性问题。

内容的提问来源于stack exchange，提问作者sam