Блог

Я хотел бы создать пользовательский слой, tensorflow-probability который я затем могу использовать для создания задней части DenseVariational слоя.

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

 def posterior_trainable_bijector(kernel_size, bias_size=0, dtype=None):
    n = kernel_size   bias_size
    c = np.log(np.expm1(1.0))

    return tf.keras.Sequential(
        [
            tfp.layers.VariableLayer(2 * n, dtype=dtype),
            tfp.layers.DistributionLambda(
                lambda t: tfp.distributions.TransformedDistribution(
                    tfd.Independent(
                        tfd.Normal(loc=tf.zeros(n), scale=tf.ones(n)),
                        reinterpreted_batch_ndims=1,
                    ),
                    tfp.bijectors.Chain(
                        bijectors=[
                            tfp.bijectors.Shift(t[..., n:]),
                            tfp.bijectors.Scale(
                                1e-5   0.01 * tf.math.softplus(c   t[..., :n])
                            ),
                        ]
                    ),
                )
            ),
        ]
    )

 

В качестве следующего шага я подумал, что было бы неплохо подклассировать DistributionLambda слой, потому что это позволило бы создать более сложные биекторы.
К сожалению, мой первый набросок, похоже, не работает. Чтобы быть более конкретным, я все еще могу запускать свой код, но он выглядит так, как будто loc_params / scale_params не обновляется во время обучения, но я не мог понять, почему это так.
Есть какие-нибудь предложения?

 class LocScaleBijectorLayer(tfp.layers.DistributionLambda):
    def __init__(
        self,
        event_shape=(),
        convert_to_tensor_fn=tfd.Distribution.sample,
        validate_args=False,
        name="LocScaleBijectorLayer",
        **kwargs,
    ):

        c = np.log(np.expm1(1.0))
        with tf.name_scope(name) as name:
            loc_params = tf.Variable(
                tf.zeros(event_shape), name="loc_var", trainable=True
            )
            scale_params = tf.Variable(
                tf.ones(event_shape), name="scale_var", trainable=True
            )

            self.base_distribution = tfd.Independent(
                tfd.Normal(loc=tf.zeros(event_shape), scale=tf.ones(event_shape)),
                reinterpreted_batch_ndims=-1,
            )

            self.bijector = tfp.bijectors.Chain(
                bijectors=[
                    tfp.bijectors.Shift(loc_params),
                    tfp.bijectors.Scale(
                        1e-5   0.01 * tf.math.softplus(c   scale_params)
                    ),
                ]
            )

            super(LocScaleBijectorLayer, self).__init__(
                lambda t: tfp.distributions.TransformedDistribution(
                    self.base_distribution, self.bijector
                ),
                convert_to_tensor_fn,
                name=name,
                **kwargs,
            )