如何在 convert_to_tensor_fn 中为 TensorFlow Probability DistributionLambda 使用来自先前 Keras 层的值答案

【问题标题】：How to use values from previous Keras layer in convert_to_tensor_fn for TensorFlow Probability DistributionLambda如何在 convert_to_tensor_fn 中为 TensorFlow Probability DistributionLambda 使用来自先前 Keras 层的值
【发布时间】：2022-10-24 21:29:37
【问题描述】：

我有一个 Keras/TensorFlow 概率模型，我想在下面的 DistributionLambda 层的 convert_to_tensor_fn 参数中包含来自前一层的值。理想情况下，我希望我能做这样的事情：

from functools import partial
import tensorflow as tf
from tensorflow.keras import layers, Model
import tensorflow_probability as tfp
from typing import Union
tfd = tfp.distributions

zero_buffer = 1e-5


def quantile(s: tfd.Distribution, q: Union[tf.Tensor, float]) -> Union[tf.Tensor, float]:
    return s.quantile(q)


# 4 records (1st value represents CDF value, 
#            2nd represents location, 
#            3rd represents scale)
sample_input = tf.constant([[0.25, 0.0, 1.0], 
                            [0.5, 1.0, 0.5], 
                            [0.75, -1.0, 2.0], 
                            [0.95, 3.0, 2.5]], dtype=tf.float32)

# Build toy model for demonstration
input_layer = layers.Input(3)
dist = tfp.layers.DistributionLambda(
    make_distribution_fn=lambda t: tfd.Normal(loc=t[..., 1],
                                              scale=zero_buffer + tf.nn.softplus(t[..., 2])),
    convert_to_tensor_fn=lambda t, s: partial(quantile, q=t[..., 0])(s)
)(input_layer)
model = Model(input_layer, dist)

但是，根据the documentation，convert_to_tensor_fn 只需要一个tfd.Distribution 作为输入； convert_to_tensor_fn=lambda t, s: 代码在上面的代码中不起作用。

如何从convert_to_tensor_fn 中的前一层访问数据？我假设有一种聪明的方法可以创建一个 partial 函数或类似的函数来让它工作。

在 Keras 模型框架之外，使用类似于以下示例的代码很容易做到这一点：

# input data in Tensor Constant form
cdf_data = tf.constant([0.25, 0.5, 0.75, 0.95], dtype=tf.float32)
norm_mu = tf.constant([0.0, 1.0, -1.0, 3.0], dtype=tf.float32)
norm_scale = tf.constant([1.0, 0.5, 2.0, 2.5], dtype=tf.float32)

quant = partial(quantile, q=cdf_data)
norm = tfd.Normal(loc=norm_mu, scale=norm_scale)
quant(norm)

输出：

<tf.Tensor: shape=(4,), dtype=float32, numpy=array([-0.6744898,  1.       ,  0.3489796,  7.112134 ], dtype=float32)>

【问题讨论】：

标签： python tensorflow keras tensorflow-probability

【解决方案1】：

我自己找到了解决此问题的方法，并决定将其发布在这里。

您可以为tfp.Normal 分发创建一个包装器类，它接受cdf 值作为参数，然后覆盖几个方法来执行您想要的操作。您尤其需要覆盖_sample_n 方法并将其替换为分位数函数，而不是从分布中随机抽取。该类看起来像这样：

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.internal import dtype_util, tensor_util, reparameterization, samplers
from tensorflow_probability.python.internal import prefer_static as ps
tfd = tfp.distributions


class NormalWrapper(tfp.distributions.Normal):
    def __init__(self,
                 loc,
                 scale,
                 cdf_vals,
                 validate_args=False,
                 allow_nan_stats=True,
                 name='NormalCDF'):
        parameters = dict(locals())
        with tf.name_scope(name) as name:
            dtype = dtype_util.common_dtype([loc, scale], dtype_hint=tf.float32)
            self._cdf_vals = tensor_util.convert_nonref_to_tensor(
                cdf_vals, dtype=dtype, name='cdf_vals')
        super(NormalWrapper, self).__init__(loc=loc,
                                            scale=scale,
                                            validate_args=validate_args,
                                            allow_nan_stats=allow_nan_stats,
                                            name=name)
        self._parameters = parameters

    def _parameter_properties(self, dtype=tf.float32, num_classes=None):
        return dict(
            loc=tfp.util.ParameterProperties(),
            scale=tfp.util.ParameterProperties(
                default_constraining_bijector_fn=(
                    lambda: tf.nn.softplus(low=dtype_util.eps(dtype)))),
            cdf_vals=tfp.util.ParameterProperties(),
        )

    @property
    def cdf_vals(self):
        return self._cdf_vals

    def _sample_n(self, n, seed=None):
        loc = tf.convert_to_tensor(self.loc)
        scale = tf.convert_to_tensor(self.scale)
        cdf_vals = tf.convert_to_tensor(self.cdf_vals)
        shape = ps.concat([[n], self._batch_shape_tensor(loc=loc, scale=scale, cdf_vals=cdf_vals)], axis=0)
        return tf.reshape(self.quantile(cdf_vals), shape=shape)

一旦你有了那个类，你可以像这样创建你的DistributionLambda 层：

dist = tfp.layers.DistributionLambda(
    make_distribution_fn=lambda t: NormalWrapper(loc=t[..., 1],
                                                 scale=zero_buffer + tf.nn.softplus(t[..., 2]),
                                                 cdf_vals=t[..., 0]),
)(input_layer)

【讨论】：

【解决方案2】：

分位数 Fn 用于提高结果的性能。它影响数据的呈现和计算。

他们在学习过程中精确的学习样本，但由于仪器或学生的不同而隔离结果。

示例：Qualaine 和正态分布

创建一个与模型一起使用的 DistributionLamda 层，然后对表示层的结果进行 Qualatine。

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

from typing import Union

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Functions
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""   
def normal_sp(params):
    return tfd.Normal(loc=params,
                      scale=1e-5 + 0.00001*tf.keras.backend.exp(params))# both parameters are learnable
                      
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Model
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""   
layer_0 = tf.keras.layers.Dense(32, activation='relu')
result_0 = layer_0( tf.constant([0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.], shape=(1, 11)) )

layer_1 = tfp.layers.DistributionLambda( normal_sp )

# Get quartiles of x with various interpolation choices.
x = tf.constant([0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.], shape=(1, 11))
                      
model = tf.keras.Sequential([
    tf.keras.Input(shape=(11)),
    layer_0,
    layer_1,
])

model.summary()

result = model.predict(x)
print( tfp.stats.quantiles(result, num_quantiles=4, interpolation='nearest') )

输出：

Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 dense (Dense)               (None, 32)                384

 distribution_lambda (Distri  ((None, 32),             0
 butionLambda)                (None, 32))

=================================================================
Total params: 384
Trainable params: 384
Non-trainable params: 0
_________________________________________________________________
1/1 [==============================] - 0s 110ms/step
tf.Tensor(
[-3.3217777e-05  1.1406353e-05  5.9894159e-05  3.1857936e+00
  9.6232319e+00], shape=(5,), dtype=float32)

【讨论】：