在训练相同的Keras模型时，如何打印模型各层的激活？

Question

以下代码（基于https://keras.io/getting-started/faq/#how-can-i-obtain-the-output-of-an-intermediate-layer）使我们能够打印模型中间层的输出，并给出一些输入，我们需要提供这些输入。具体来说，在此示例中，鉴于层dense的输入，我正在打印层input_1的输出，该输入也恰好是我模型的输入层（但这不必这样）。

import numpy as np
import tensorflow as tf


def get_model():
    inp = tf.keras.layers.Input(shape=(1,))
    x = tf.keras.layers.Dense(8)(inp)
    x = tf.keras.layers.Dense(16)(x)
    out = tf.keras.layers.Dense(1)(x)
    model = tf.keras.Model(inputs=inp, outputs=out)
    model.summary()
    return model


def train():
    my_model = get_model()
    my_model.compile(optimizer="adam", loss="mse")

    get_layer_output = tf.keras.backend.function([my_model.get_layer("input_1").input],
                                                 [my_model.get_layer("dense").output])

    data_x = np.array([[1], [2], [3], [4]])
    layer_output = get_layer_output(data_x)[0]
    print(layer_output)


if __name__ == '__main__':
    train()

但是，我想在训练模型的同时（即在每个小批量之后）打印给定对应的上一层（由模型定义）的输出的每一层的输出。我尝试使用callback，它调用tf.print来打印模型的输出，但是我遇到了一个错误，如this Github issue中所述（即TensorFlow中有一个错误） 2.0，这是我正在使用并且要使用的版本）。

更清楚地说，我想在训练模型时调试每层的输出，以便我可以理解输入如何流过每一层以及激活程度不是太高（爆炸）还是太小（消失）。我可以迭代地向get_layer_output提供一批数据，但是我希望能够在使用fit或fit_generator训练模型的同时打印每层的激活。此外，我想了解激活值如何从输入层演变为输出层，而不仅仅是在给定输入到另一层的情况下打印一层的激活。

这里的a Github issue提出类似的要求。