将可训练双射器嵌入 Keras 模型

Question

我正在尝试实现嵌入在 Keras 模型中的规范化流程。在我能找到的所有示例中，例如 MAF 的文档中，构成规范化流的双射器都嵌入到 TransformedDistribution 中并直接用于训练等。

我正在尝试将此 TransformedDistribution 嵌入到 keras 模型中，以匹配我拥有的其他模型的架构，这些模型继承自 keras 模型。

不幸的是，到目前为止，我所有的尝试（参见代码）都未能将转换后的分布中的可训练变量转移到 keras 模型。

我试图让双射器从tf.keras.layers.Layer 继承，这并没有改变任何东西。

import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions
tfb = tfp.bijectors


class Flow(tfb.Bijector, tf.Module):
    """
    tf.Module to register trainable_variables
    """

    def __init__(self, d, init_sigma=0.1, **kwargs):
        super(Flow, self).__init__(
            dtype=tf.float32,
            forward_min_event_ndims=0,
            inverse_min_event_ndims=0,
            **kwargs
        )
        # Shape of the flow goes from Rd to Rd
        self.d = d
        # Weights/Variables initializer
        self.init_sigma = init_sigma
        w_init = tf.random_normal_initializer(stddev=self.init_sigma)
        # Variables
        self.u = tf.Variable(
            w_init(shape=[1, self.d], dtype=tf.float32),
            dtype=tf.float32,
            name='u',
            trainable=True,
        )

    def _forward(self, x):
        return x

    def _inverse(self, y):
        return y


class Flows(tf.keras.Model):

    def __init__(self, d=2, shape=(100, 2), n_flows=10, ):
        super(Flows, self).__init__()
        # Parameters
        self.d = d
        self.shape = shape
        self.n_flows = n_flows
        # Base distribution - MF = Multivariate normal diag
        base_distribution = tfd.MultivariateNormalDiag(
            loc=tf.zeros(shape=shape, dtype=tf.float32)
        )
        # Flows as chain of bijector
        flows = []
        for n in range(n_flows):
            flows.append(Flow(self.d, name=f"flow_{n + 1}"))
        bijector = tfb.Chain(list(reversed(flows)))
        self.flow = tfd.TransformedDistribution(
            distribution=base_distribution,
            bijector=bijector
        )

    def call(self, *inputs):
        return self.flow.bijector.forward(*inputs)

    def log_prob(self, *inputs):
        return self.flow.log_prob(*inputs)

    def sample(self, num):
        return self.flow.sample(num)


q = Flows()
# Call to instantiate variables
q(tf.zeros(q.shape))
# Prints no trainable params
print(q.summary())
# Prints expected trainable params
print(q.flow.trainable_variables)

知道这是否可能吗？谢谢！

Answer 1

我也遇到了这个问题。这似乎是由 TFP 和 TF 2.0 之间的不兼容问题引起的（一些相关问题 https://github.com/tensorflow/probability/issues/355 和 https://github.com/tensorflow/probability/issues/946）。

作为一种解决方法，您需要将转换后的分布/双射器的（可训练）变量作为属性添加到您的 Keras 模型中：

class Flows(tf.keras.Model):

    def __init__(self, d=2, shape=(100, 2), n_flows=10, ):
        super(Flows, self).__init__()
        # Parameters
        self.d = d
        self.shape = shape
        self.n_flows = n_flows
        # Base distribution - MF = Multivariate normal diag
        base_distribution = tfd.MultivariateNormalDiag(
            loc=tf.zeros(shape=shape, dtype=tf.float32)
        )
        # Flows as chain of bijector
        flows = []
        for n in range(n_flows):
            flows.append(Flow(self.d, name=f"flow_{n + 1}"))
        bijector = tfb.Chain(list(reversed(flows)))
        self.flow = tfd.TransformedDistribution(
            distribution=base_distribution,
            bijector=bijector
        )
        # issue: https://github.com/tensorflow/probability/issues/355, https://github.com/tensorflow/probability/issues/946
        # need to add bijector's trainable variables as an attribute (name does not matter)
        # otherwise this layer has zero trainable variables
        self._variables = self.flow.variables # https://github.com/tensorflow/probability/issues/355

    def call(self, *inputs):
        return self.flow.bijector.forward(*inputs)

    def log_prob(self, *inputs):
        return self.flow.log_prob(*inputs)

    def sample(self, num):
        return self.flow.sample(num)

添加后，您的模型应该有可训练的变量和权重进行优化。