TensorFlow Probability循环构建网络仅输出最后一层分布的问题

问题根源 #

这是Python闭包延迟绑定导致的坑。在循环里创建DistributionLambda时，你写的匿名函数并没有捕获当前迭代的变量值，而是拿着变量的引用。等模型真正执行的时候，循环早就跑完了，所有闭包都会指向循环最后一次迭代的变量值，自然所有输出分布都变成最后一个了。

至于你说移除total = 0会崩溃，也是因为循环里的变量（比如total）被闭包引用，循环结束后如果没重新赋值，会触发变量作用域的冲突或未定义问题。

修复方法 #

方法1：用默认参数捕获当前循环变量 #

在创建DistributionLambda的闭包时，把当前的分布类作为默认参数传入，这样每次迭代都会把当前值固定下来：

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

# 示例分布映射字典
dist_lookup = {
    'normal': tfd.Normal,
    'exponential': tfd.Exponential
}

# 基础模型结构
inputs = tf.keras.Input(shape=(10,))
x = tf.keras.layers.Dense(32, activation='relu')(inputs)

outputs = []
dist_names = ['normal', 'exponential']  # 动态分布列表
total = 0

for dist_name in dist_names:
    dist_cls = dist_lookup[dist_name]
    # 用默认参数dist=dist_cls捕获当前迭代的分布类
    output_layer = tfp.layers.DistributionLambda(
        lambda params, dist=dist_cls: dist(params[0], params[1])
    )(x)
    outputs.append(output_layer)
    total += 1

model = tf.keras.Model(inputs=inputs, outputs=outputs)

方法2：用辅助函数生成分布层 #

写个小函数专门生成对应分布的DistributionLambda，每次调用都会创建新的作用域，把当前分布类牢牢锁住：

def make_dist_layer(dist_cls):
    return tfp.layers.DistributionLambda(
        lambda params: dist_cls(params[0], params[1])
    )

# 循环里调用辅助函数
outputs = []
for dist_name in dist_names:
    dist_cls = dist_lookup[dist_name]
    output_layer = make_dist_layer(dist_cls)(x)
    outputs.append(output_layer)

方法3：用functools.partial绑定参数 #

借助functools.partial把分布类直接绑定到lambda函数上，避开闭包引用的坑：

from functools import partial

outputs = []
for dist_name in dist_names:
    dist_cls = dist_lookup[dist_name]
    dist_fn = partial(lambda dist, params: dist(params[0], params[1]), dist_cls)
    output_layer = tfp.layers.DistributionLambda(dist_fn)(x)
    outputs.append(output_layer)

验证效果 #

修复后再跑预测、画直方图，每个输出都会对应你定义的分布类型，不会再全部复用最后一个分布的特征。

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