在TensorFlow模型中的隐藏层之间打印值

import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Dense(units=64, activation='relu', input_shape=(input_dim,)),
    tf.keras.layers.Dense(units=32, activation='relu'),
    tf.keras.layers.Dense(units=output_dim, activation='softmax')
])

class HiddenLayerPrintCallback(tf.keras.callbacks.Callback):
    def __init__(self, layer_index):
        super(HiddenLayerPrintCallback, self).__init__()
        self.layer_index = layer_index

    def on_epoch_end(self, epoch, logs=None):
        hidden_layer_model = tf.keras.Model(inputs=self.model.inputs,
                                            outputs=self.model.layers[self.layer_index].output)
        hidden_layer_output = hidden_layer_model.predict(self.model.inputs)
        print("Hidden layer {} output: {}".format(self.layer_index, hidden_layer_output))

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=10, callbacks=[HiddenLayerPrintCallback(layer_index=1)])