使用TF 2

Question

在使用model.fit(…)进行训练的同时，是否可以在tf.keras模型的内记录不同层的输出（又称为激活）？

入门：

具有类似的模型

inp = tf.keras.layers.Input((10,))
state = Dense(10)(inp)
outp = Dense(10)(state)
model = tf.keras.Model(inp, outp)

我想在state过程中创建流过model.fit(…)的值的直方图。

不满意的解决方案：

以下代码可用于绘制直方图：

# -*- coding: utf-8 -*-
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense


# Summary Layer
class Layer(tf.keras.layers.Layer):
    def __init__(self, writer):
        super(Layer, self).__init__()
        self.writer = writer

    def call(self, input):
        with self.writer.as_default():
            tf.summary.histogram("hist", input, step=0)
        return input


# Data
x, y = np.random.rand(10000, 10), np.random.rand(10000, 10)
writer = tf.summary.create_file_writer("logs")

# Model
inp = tf.keras.layers.Input((10,))
state = Dense(10)(inp)
state = Layer(writer)(state)
outp = Dense(10)(state)
model = tf.keras.Model(inp, outp)
model.compile("sgd", "mse")

# Fit
model.fit(x, y, epochs=3)

但是，这里step的{​​{1}}参数始终设置为0。 这样，所有图都在同一步骤中绘制。

问题：

由于tf.summary.histogram已经过tf 2，所以我看不到解决方案，因此在设置global_step时如何知道步长值。 是否有针对此问题的解决方法或现成的解决方案？

非常感谢您的帮助。 祝一切顺利！

附录：

在阅读文档时，我发现类似：

tf.summary.histogram

当然，步长值是已知的。 但是在训练中使用后一个代码段记录模型层的输出有两个问题：

1. 拟合后再次
需要评估模型中的每个张量
2. 其他值在拟合后流过张量

对我来说，解决方案似乎非常不自然，我想我可能正在监督某些事情。