自定义 tf.keras.keras.callbacks.TensorBoard 在 tensorflow 版本 >= 1.15.0 中无法正常工作

Customized tf.keras.keras.callbacks.TensorBoard does not work well in tensorflow version >= 1.15.0

以下定制的 TensorBoard,从 this github repo 克隆而来,旨在在每批结束时存储学习率和所谓的 KL 权重,在 tensorflow 版本 <= 1.12.0 中运行良好,但是版本 >= 1.15.0 中的每个批次都没有做它应该做的事情。我该如何解决?

class TensorBoardLR(TensorBoard):
    """ A modification to the Tensorboard callback to also include the scalars of learning rate and KL weight"""
    def __init__(self, *args, **kwargs):
        self.kl_weight = kwargs.pop('kl_weight')
        super().__init__(*args, **kwargs)
        self.count = 0

    def on_batch_end(self, batch, logs=None):
        logs.update({'lr': K.eval(self.model.optimizer.lr),
                     'kl_weight': K.eval(self.kl_weight)})
        super().on_batch_end(batch, logs)

将您的方法名称从 on_batch_end 更改为 on_train_batch_end

他们保留了遗留代码的方法,https://www.tensorflow.org/versions/r1.15/api_docs/python/tf/keras/callbacks/Callback这也是keras和tensorflow.keras的区别。