3

我正在尝试使用 tf.keras 重构我们的一些强化学习方法。keras 模型似乎非常适合我用于价值函数和预处理器的前馈网络，但是在尝试使用 keras 实现一些概率模型（例如策略）时，我遇到了困难。特别是，将 tfp.distributions 和 tfp.bijectors 与 tf.keras.Model 结合起来对我来说感觉非常不直观，而且我看到的所有示例（例如 [1, 2, 3]）要么过于简单，要么依赖在我看来，“黑客”似乎破坏了使用 keras 模型的许多好处（例如，能够从模型本身隐藏输入处理、会话和数值评估）。

假设我想实现一个潜在空间策略，如 [4] 中所述，它使用 RealNVP 流将基于状态的高斯样本转换为动作。该策略需要至少支持两种操作： 1. 对动作 Y 进行采样，使得 Y = g(X|S) X ~ Normal(0, 1)，其中 g 是 [4, 5] 中描述的 RealNVP 变换， S 是条件变量（例如，RL 案例中的状态观察）。2. 计算采样 Y 的对数概率。

一个简单的实现可能如下所示：

class LearnableConditionalRealNVP(object):
    def __init__(self, input_shape, output_shape):
        self._input_shape = input_shape
        self._output_size = np.prod(output_shape)

        conditions = tf.keras.layers.Input(shape=input_shape)

        batch_size = tf.keras.layers.Lambda(
            lambda x: tf.shape(x)[0])(conditions)

        def samples_and_log_probs_fn(inputs):
            conditions, batch_size = inputs

            base_distribution = tfp.distributions.MultivariateNormalDiag(
                loc=tf.zeros(output_shape),
                scale_diag=tf.ones(output_shape))

            real_nvp_bijector = tfp.bijectors.RealNVP(
                num_masked=self._output_size // 2,
                shift_and_log_scale_fn=conditioned_real_nvp_template(
                    hidden_layers=(128, 128),
                    activation=tf.nn.relu),
                name='real_nvp')

            distribution = (
                tfp.distributions.ConditionalTransformedDistribution(
                    distribution=base_distribution,
                    bijector=real_nvp_bijector))

            samples = distribution.sample(batch_size)
            log_probs = distribution.log_prob(samples)

            return [samples, tf.reshape(log_probs, (-1, 1))]

        samples, log_probs = tf.keras.layers.Lambda(
            samples_and_log_probs_fn)([conditions, batch_size])

        self.samples_and_log_probs_model = tf.keras.Model(
            conditions, [samples, log_probs])

    def samples_and_log_probs(self, conditions):
        return self.samples_and_log_probs_model(conditions)

    def samples_and_log_probs_np(self, conditions):
        return self.samples_and_log_probs_model.predict(conditions)

其中conditioned_real_nvp_template创建了一个前馈网络，它将潜在样本和沿最后一个轴的条件值连接起来，并将它们用作其输入。完整的例子可以在这里找到：https ://gist.github.com/hartikainen/17ac2ec102032e986cb4d31e225f592a

这种处理分配的方式对我来说有两个主要好处。首先，我不必手动处理参数的重用。我可以samples_and_log_probs在我的代码中多次调用，它会自动重用模型的参数。其次，如果我想获得数字输出，我不需要知道任何关于会话的信息。将中间层包装到它们自己的模型中允许我调用处理会话的预测方法。

进一步扩展示例时会出现问题。假设我想修改 ，LearnableConditionalRealNVP以便我可以提供潜在样本 x 作为输入，而不是调用distribution.sample()，samples_and_log_probs_fn而是返回distribution.forward(x)。或者，也许我想将样本和 log_probs 分别从模型中取出。这将需要我将其拆分samples_and_log_probs_fn为两个单独的 lambda 函数，但如果我想共享 RealNVP 双射器的参数，这样做并非易事（因为我不能将双射器作为输入/输出传递给/来自 keras 层）。

LearnableConditionalRealNVP我尝试通过子类化from 来解决这些问题tf.keras.Model，但我所有的尝试都导致实现混乱，主要是由于输入和输出的变化。具体来说，我无法call为模型创建 - 方法以使模型保留与 一起使用的能力predict，并且我必须在 - 方法中制作一些技巧__call__。这些都不是很可怕，但它们确实增加了使用 keras 模型的开销，以至于我更容易在普通的 tensorflow 中实现这些类型的东西并手动处理会话、numpy 输出等。

我的问题是：

  1. 张量流双射器/分布是否应该与 keras（模型）兼容？如果是，有谁知道是否有任何重要的例子可供我看一下？如果不是，是否有计划使它们兼容？
  2. 在我认为一个模型有多个不同输出的情况下，我应该如何使用 keras 模型？例如，在我上面的例子中，可学习的 RealNVP 分布直观地感觉它应该是一个模型，但它有多个，可能是独立的输入/输出，这使得它很难适应 keras 模型框架。我愿意接受我的直觉是错误的，在这种情况下，很高兴听到构建此类模型的最佳实践是什么。
  3. 有没有办法像 [1] 中所做的那样将非张量数据作为输入/输出传递给 keras 模型，同时仍然保持模型以某种方式连接。如果在此示例中使用输入而不是tfe.Variables，则它会中断，因为图形未连接。

编辑：在发布此内容并使用上述实现进行一些额外测试后，我注意到该模型毕竟不可训练，因为 RealNVP 双射器的变量是在 keras lambda 层中创建的。这表明我构建这些模型的功能方式根本不能用于这些类型的模型。

[1] https://github.com/tensorflow/probability/blob/5f5510201865350b6cce2a0f18fbe0cdf4f15eee/tensorflow_probability/examples/disentangled_vae.py#L186

[2] https://blog.keras.io/building-autoencoders-in-keras.html

[3] http://louistiao.me/posts/implementing-variational-autoencoders-in-keras-beyond-the-quickstart-tutorial/

[4] https://arxiv.org/pdf/1804.02808.pdf

[5] https://arxiv.org/abs/1605.08803

4

0 回答 0