Извлечение изученных NN параметров распределения задней массы из слоя DenseVariational

Question

Я также разместил этот вопрос в выпусках GitHub с вероятностью тензорного потока: https://github.com/tensorflow/probability/issues/892

Я использую Tensorflow 2.1.0 и вероятность тензорного потока 0,9,0 в python 3,6 +0,8. Я работаю с моделью Tensorflow Probability Keras, в которой слой DenseVariational определен следующим образом (взято из примеров, найденных в Интернете):

def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    c = np.log(np.expm1(1.))
    return tf.keras.Sequential([
        tfp.layers.VariableLayer(2 * n, dtype=dtype),
        tfp.layers.DistributionLambda(lambda t: tfd.Independent(
            tfd.Normal(loc=t[..., :n], scale=1e-5 + tf.nn.softplus(c + t[..., n:])),
            reinterpreted_batch_ndims=1)),
    ])


def prior_trainable(kernel_size, bias_size=0, dtype=None):
    n = kernel_size + bias_size
    return tf.keras.Sequential([
        tfp.layers.VariableLayer(n, dtype=dtype),
        tfp.layers.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t, scale=1),
                                                                reinterpreted_batch_ndims=1)),
    ])

dense = tfp.layers.DenseVariational(units=units, make_posterior_fn=posterior_mean_field,
                                                             make_prior_fn=prior_trainable,
                                                            )(prev_layer)

Если я обучу свою модель и затем удаляю слои, следующие за этим слоем, Оставшаяся модель выведет случайные величины из изученных задних распределений веса. Примерно так:

from tensorflow.keras import Model
# DenseVariational layer is 3rd to last layer in this case
cropped_model = Model(inputs, model.layers[-3].output)  
cropped_mode.predict(test_data)

В большинстве случаев это нормально (например, тренировка, выборка и т. Д. c.). Однако существует ли прямой способ получить усвоенные значения lo c и апостериорные значения, возвращаемые для данного ввода (например, test_data) в эту cropped_model, вместо выборочного извлечения из определенного ими распределения?