在 tensorflow 中实现 KL 预热：回调中的 tf.keras.backend.variable 在 epoch 上不稳定

Question

我正在尝试在 TensorFlow 中实现带有 KL 预热的变分自动编码器的变体（论文 here）。这个想法是，损失的 KL 项应该在训练开始时在指定数量的 epoch 上线性增加。

我尝试的方法是使用一个回调，在每次新纪元开始时在 K.variable 中设置一个值，作为预热所需跨度内的当前纪元数（例如，如果预热设置为 last对于 10 个时期，在第 6 个时期，损失中的 KL 项应乘以 0.6）。

我还在 KL 中包含一个 add_metric()（它作为层 subclass 实现）以在训练期间控制 kl_rate。问题是变量的值不稳定！它在每个新纪元开始时接近期望值，但在每次迭代中都会缓慢衰减，使得过程不太可控。

你知道我做错了什么吗？我也不确定这是回调本身的问题（以及随后的实际使用值）还是报告的指标。

谢谢！

进口：

import tensorflow.keras.backend as K

回调（self.kl_warmup是模型class的一个参数，设置为整数，对应应该增加kl率的epoch数）：

kl_beta = K.variable(1.0, name="kl_beta")
if self.kl_warmup:

    kl_warmup_callback = LambdaCallback(
        on_epoch_begin=lambda epoch, logs: K.set_value(
            kl_beta, K.min([epoch / self.kl_warmup, 1])
        )
    )

z_mean, z_log_sigma = KLDivergenceLayer(beta=kl_beta)([z_mean, z_log_sigma])

KL层：

class KLDivergenceLayer(Layer):

""" Identity transform layer that adds KL divergence
to the final model loss.
"""

def __init__(self, beta=1.0, *args, **kwargs):
    self.is_placeholder = True
    self.beta = beta
    super(KLDivergenceLayer, self).__init__(*args, **kwargs)

def get_config(self):
    config = super().get_config().copy()
    config.update({"beta": self.beta})
    return config

def call(self, inputs, **kwargs):
    mu, log_var = inputs
    kL_batch = -0.5 * K.sum(1 + log_var - K.square(mu) - K.exp(log_var), axis=-1)

    self.add_loss(self.beta * K.mean(kL_batch), inputs=inputs)
    self.add_metric(self.beta, aggregation="mean", name="kl_rate")

    return inputs

模型实例（整个模型构建在 class 中，returns 编码器、生成器、完整的 vae 和 kl_rate 回调：

encoder, generator, vae, kl_warmup_callback = SEQ_2_SEQ_VAE(pttest.shape,
                                                               loss='ELBO',
                                                               kl_warmup_epochs=10).build()

fit() 调用：

history = vae.fit(x=pttrain, y=pttrain, epochs=100, batch_size=512, verbose=1,
              validation_data=(pttest, pttest),
              callbacks=[tensorboard_callback, kl_warmup_callback])

训练过程的一个片段（注意 kl_rate 应该是零并且它是关闭的）：

来自 tensorboard 的 kl_rate 个时期的屏幕截图（跨度设置为 10 个时期；在 10 个时期之后它应该达到 1，但它收敛到大约 0.9）

Answer 1

经过更多研究，我最终自己发现了它。

kl_beta._trainable = False

成功了:) 谢谢！

Answer 2

使用回调可以，但有点笨拙。如果您对训练步骤 (iterations) 而不是 epochs 感到满意，那么该值在优化器中已经是 tracked/updated。

如果新层没有任何变量或执行任何操作，那么创建一个新层也有点误导。在生成分布参数的层上使用自定义 activity_regularizer 更适合这里，并且可以防止您意外使用未正则化的参数。

没有见到你 SEQ_2_SEQ_VAE 很难给出一个准确的代码示例，但希望下面的内容足以说明如何实施。

class KLDivergenceRegularizer(tf.keras.regularizers.Regularizer):
  def __init__(self, iters: tf.Variable, warm_up_iters: int):
    self._iters = iters
    self._warm_up_iters = warm_up_iters

  def __call__(self, activation):
    # note: activity regularizers automatically divide by batch size
    mu, log_var = activation
    k = K.min(self._iters / self._warm_up_iters, 1)
    return -0.5 * k * K.sum(1 + log_var - K.square(mu) - K.exp(log_var))

optimizer = tf.keras.optimizers.Adam()  # or whatever optimizer you want
warm_up_iters = 1000  # not epochs

inp = make_input()
x = tf.keras.layers.Dense(...)(inp)
...
mu, log_sigma = ParametersLayer(
  ...,
  activity_regularizer=KLDivergenceRegularizer(optimizer.iterations, warm_up_iters))(x)

...
vae = bulid_vae(...)
vae.compile(optimizer=optimizer, ...)

如果你有多个这样的东西，你可以在正则化器中使用 tf.summary.experimental.set_step(optimizer.iterations) 和 tf.summary.experimental.get_step()。

在 tensorflow 中实现 KL 预热：回调中的 tf.keras.backend.variable 在 epoch 上不稳定

Implementing KL warmup in tensorflow: tf.keras.backend.variable in callback is unstable over epochs

python

callback

deep-learning

keras

tensorflow2.0