训练预测后batch_size显示为None的原因及解决方案求助

你的问题根源在DistributionLambda的采样逻辑上：调用s.sample(5)时，TensorFlow会将采样次数作为第一个维度，导致输出形状为(5, None, 1)（None是动态推断的batch维度）。当部分数据的batch维度无法被静态确定时，就会出现你看到的[15, None, 1]（实际是维度顺序颠倒，且batch维度被标记为动态）。

以下是几种可行的解决方法：

方法1：调整采样维度顺序，确保batch维度在前 #

修改DistributionLambda的convert_to_tensor_fn，将采样后的维度转置，让batch维度处于第一位：

output = tfpl.DistributionLambda(
        lambda t: tfd.Normal(loc=t, scale=1),
        convert_to_tensor_fn=lambda s: tf.transpose(s.sample(5), perm=[1, 0, 2])
)(x)

修改后，模型输出形状为(None, 5, 1)，预测时predt.shape会变为(15, 5, 1)，batch维度明确无None。

方法2：移除转换函数，预测时手动采样（更灵活） #

让模型直接输出分布对象，在预测阶段按需采样，这种方式能避免动态维度的歧义：

# 移除convert_to_tensor_fn参数，模型输出分布对象
output = tfpl.DistributionLambda(
        lambda t: tfd.Normal(loc=t, scale=1)
)(x)

# 训练流程保持不变
model.compile(loss=negloglik, optimizer=optimizer)
model.fit(X_train, y_train, epochs=300, verbose=1)

# 预测时手动采样并调整维度顺序
pred_dist = model.predict(X_train, verbose=0)
predt = pred_dist.sample(5)
predt = tf.transpose(predt, perm=[1, 0, 2])
print(predt.shape)  # 输出形状为(15, 5, 1)

方法3：固定输入batch_size（不推荐） #

如果你的训练和预测始终使用固定大小的batch，可以在输入层直接指定batch_size：

input = tf.keras.Input((6,1), batch_size=15, name='input')

但这种方式灵活性差，无法适配不同batch大小的场景，仅作为备选方案。

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