如何使用“tf.keras.callbacks.TensorBoard”记录激活？

Question

的文档tf.keras.callbacks.TensorBoard说明该工具可以做到这一点：

此回调记录 TensorBoard 的事件，包括：

• 指标汇总图
• 训练图可视化
• 激活直方图
• 抽样分析

稍后还有：

histogram_freq：计算模型各层的激活和权重直方图的频率（以纪元为单位） 。如果设置为 0，则不会计算直方图。必须为直方图可视化指定验证数据（或分割）。

但是，当使用此参数时，我没有看到任何写入的激活摘要，只写入了权重本身。查看源代码，我也没有看到任何与激活相关的内容。

那么我错过了什么吗？是否可以在 TF2 中编写没有自定义代码的激活摘要？

Answer 1

根据我对创建 keras 回调以在训练期间保存每批的模型预测和目标的回答，我使用以下代码：

"""Demonstrate activation histograms."""
import tensorflow as tf
from tensorflow import keras


class ActivationHistogramCallback(keras.callbacks.Callback):
    """Output activation histograms."""

    def __init__(self, layers):
        """Initialize layer data."""
        super().__init__()

        self.layers = layers
        self.batch_layer_outputs = {}
        self.writer = tf.summary.create_file_writer("activations")
        self.step = tf.Variable(0, dtype=tf.int64)

    def set_model(self, _model):
        """Wrap layer calls to access layer activations."""
        for layer in self.layers:
            self.batch_layer_outputs[layer] = tf_nan(layer.output.dtype)

            def outer_call(inputs, layer=layer, layer_call=layer.call):
                outputs = layer_call(inputs)
                self.batch_layer_outputs[layer].assign(outputs)
                return outputs

            layer.call = outer_call

    def on_train_batch_end(self, _batch, _logs=None):
        """Write training batch histograms."""
        with self.writer.as_default():
            for layer, outputs in self.batch_layer_outputs.items():
                if isinstance(layer, keras.layers.InputLayer):
                    continue
                tf.summary.histogram(f"{layer.name}/output", outputs, step=self.step)

        self.step.assign_add(1)


def tf_nan(dtype):
    """Create NaN variable of proper dtype and variable shape for assign()."""
    return tf.Variable(float("nan"), dtype=dtype, shape=tf.TensorShape(None))


def main():
    """Run main."""
    model = keras.Sequential([keras.layers.Dense(1, input_shape=(2,))])

    callback = ActivationHistogramCallback(model.layers)

    model.compile(loss="mse", optimizer="adam")
    model.fit(
        x=tf.transpose(tf.range(7.0) + [[0.2], [0.4]]),
        y=tf.transpose(tf.range(7.0) + 10 + [[0.5]]),
        validation_data=(
            tf.transpose(tf.range(11.0) + 30 + [[0.6], [0.7]]),
            tf.transpose(tf.range(11.0) + 40 + [[0.9]]),
        ),
        shuffle=False,
        batch_size=3,
        epochs=2,
        verbose=0,
        callbacks=[callback],
    )


if __name__ == "__main__":
    main()

对于使用 2 epoch 和 3 个批次进行训练的示例（由于 7 个训练样本的数量为奇数，因此大小不等），然后我们会看到预期输出（6 个批次，具有 3, 3, 1, 3, 3, 1 个峰值）。

通过 200 个 epoch（600 个批次），我们还可以看到训练进度：