如何获取TensorFlow概率层模型的输出实际形状？

Hey there! I ran into this exact issue when working with TensorFlow Probability's DistributionLambda layers too—let's walk through what's going on and how to get the shape info you need.

Why the Error Happens #

When you add a DistributionLambda layer, your model's output isn't a standard Keras tensor—it's a probability distribution object wrapped in a UserRegisteredSpec. That's why calling model.output.shape throws the AttributeError: 'UserRegisteredSpec' object has no attribute '_shape' error: this spec doesn't have the same shape attributes as a regular tensor.

And when you use tf.shape(model.output), you're getting a tensor that represents the dynamic shape of the distribution's samples (computed at runtime), not the static shape definition you're looking for.

How to Get the Static Shape #

To retrieve the static shape (like [None, 1] for your sample code), you need to access the distribution's sample_shape property. Here's how to do it with your code:

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd

tfd = tfp.distributions
model = tf.keras.Sequential()
model.add(tf.keras.layers.Input(10))
model.add(tf.keras.layers.Dense(2, activation="linear"))
model.add(
    tfp.layers.DistributionLambda(
        lambda t: tfd.Normal(
            loc=t[..., :1],
            scale=1e-3 + tf.math.softplus(0.1 * t[..., 1:])
        )
    )
)

# Get static sample shape as a TensorShape object
static_shape = model.output.sample_shape
# Convert to a list (matches the format you want: [None, 1] for your example)
static_shape_list = static_shape.as_list()
print(static_shape_list)  # Output: [None, 1]

Getting Dynamic Runtime Shape #

If you need the actual shape values when running the model with real data (like knowing the batch size at runtime), you'll need to pass an input tensor first, generate a sample from the distribution, then use tf.shape() with .numpy() to get concrete values:

# Create a sample input batch (batch size 32, input dimension 10)
sample_input = tf.random.normal((32, 10))
# Get the distribution from the model
output_dist = model(sample_input)
# Generate a sample and get its dynamic shape
dynamic_shape = tf.shape(output_dist.sample())
print(dynamic_shape.numpy())  # Output: [32, 1]

Key Takeaways #

• For static shape definitions (the "inferred" shape you see in the KerasTensor description), use model.output.sample_shape.as_list().
• For runtime dynamic shapes (actual values when the model runs), pass an input batch, generate a sample from the distribution, then use tf.shape().numpy().

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