TensorFlow中Keras Layer封装TF函数后模型无训练权重问题求助

你遇到的核心问题是：将可训练分布的log_prob方法封装到自定义Keras Layer后，模型无法识别分布的训练权重，导致训练时出现“无训练权重”的警告，无法优化MaskedAutoregressiveFlow序列。

问题根源 #

你自定义的LogProb层仅直接引用了trainable_distribution，但手动赋值self._variables的方式不符合Keras的权重注册规范——Keras不会自动将这些变量纳入模型的可训练权重管理体系，因此模型编译时认为没有可优化的参数。

解决方案 #

修改自定义LogProb层，显式将分布的可训练/非训练变量注册到层的权重列表中，让Keras正确追踪这些变量。具体代码如下：

class LogProb(tfk.Layer):
    def __init__(self, trainable_distribution):
        super(LogProb, self).__init__()
        self.trainable_distribution = trainable_distribution
        
        # 显式注册分布的可训练变量
        for var in self.trainable_distribution.trainable_variables:
            self._trainable_weights.append(var)
        
        # 注册分布的非训练变量（如果有需要）
        for var in self.trainable_distribution.non_trainable_variables:
            self._non_trainable_weights.append(var)

    def call(self, x):
        return self.trainable_distribution.log_prob(x)

原理说明 #

Keras通过层的_trainable_weights和_non_trainable_weights列表来管理权重，将分布的变量直接添加到这些列表中，就能让模型识别到这些可训练参数，从而在训练时对MaskedAutoregressiveFlow的权重进行优化。

验证方法 #

修改后，你可以在模型编译前打印model.trainable_weights，应该能看到MaskedAutoregressiveFlow相关的可训练参数，训练时也不会再出现“无训练权重”的警告。

补充注意事项 #

确保你的trainable_distribution本身是带有可训练参数的——比如创建MaskedAutoregressiveFlow时设置了hidden_units等可训练配置，否则即使注册变量，也不会有可优化的参数。

内容的提问来源于stack exchange，提问作者Jean-Eric