TensorBoard 中 tf.summary.* 操作的步数始终为 0

Question

当我使用 TensorFlow 2.3 训练我的模型时，我想可视化一些中间张量，这些张量是使用我自定义 tf.keras.layers.Layer.

的计算图中的权重计算的

所以我使用 tf.summary.image() 来记录这些张量并将它们可视化为这样的图像：

class CustomizedLayer(tf.keras.layers.Layer):
    def call(self, inputs, training=None):
        # ... some code ...
        tf.summary.image(name="some_weight_map", data=some_weight_map)
        # ... some code ...

但是在TensorBoard中，无论经过多少步，都只显示一张第0步的图像。

并且我尝试将tf.summary.image()的参数step设置为从tf.summary.experimental.get_step()获取的值：

tf.summary.image(name="weight_map", data=weight_map, step=tf.summary.experimental.get_step())

并使用 tf.Variable 从自定义回调中调用 tf.summary.experimental.set_step 来更新步骤，如下所示：

class SummaryCallback(tf.keras.callbacks.Callback):
def __init__(self, step_per_epoch):
    super().__init__()
    self.global_step = tf.Variable(initial_value=0, trainable=False, name="global_step")
    self.global_epoch = 0
    self.step_per_epoch = step_per_epoch
    tf.summary.experimental.set_step(self.global_step)

def on_batch_end(self, batch, logs=None):
    self.global_step = batch + self.step_per_epoch * self.global_epoch
    tf.summary.experimental.set_step(self.global_step)  
    # whether the line above is commented, calling tf.summary.experimental.get_step() in computation graph code always returns 0.
    # tf.print(self.global_step)

def on_epoch_end(self, epoch, logs=None):
    self.global_epoch += 1

此回调的实例在 model.fit() 函数的参数 callbacks 中传递。

但是返回的值tf.summary.experimental.get_step()仍然是0。

“tf.summary.experimental.set_step()”的 TensorFlow 文档说：

when using this with @tf.functions, the step value will be captured at the time the function is traced, so changes to the step outside the function will not be reflected inside the function unless using a tf.Variable step.

根据文档，我已经使用一个变量来存储步骤，但它的变化仍然没有反映在函数内部（或keras.Model）。

注意：在我将代码迁移到 TensorFlow 2 之前，我的代码在 TensorFlow 1.x 中产生了预期的结果，只需简单的一行 tf.summary.image()。

所以我想知道我的方法在 TensorFlow 2 中是否有误？

在 TF2 中，如何在计算图中获取训练步骤？

或者有其他解决方案 在 TensorFlow 2 的模型中汇总张量（如标量、图像等）？

Answer 1

我发现此问题已在 Tensorflow 的 Github 存储库中报告：https://github.com/tensorflow/tensorflow/issues/43568

这是由于在模型中使用了tf.summary而同时启用了tf.keras.callbacks.TensorBoard回调，并且步骤将始终为零。问题报告者给出了一个临时解决方案。

要修复它，继承tf.keras.callbacks.TensorBoardclass并覆盖on_train_begin方法和 on_test_begin 这样的方法：

class TensorBoardFix(tf.keras.callbacks.TensorBoard):
"""
This fixes incorrect step values when using the TensorBoard callback with custom summary ops
"""

def on_train_begin(self, *args, **kwargs):
    super(TensorBoardFix, self).on_train_begin(*args, **kwargs)
    tf.summary.experimental.set_step(self._train_step)


def on_test_begin(self, *args, **kwargs):
    super(TensorBoardFix, self).on_test_begin(*args, **kwargs)
    tf.summary.experimental.set_step(self._val_step)

并在 model.fit():

中使用此固定回调 class

tensorboard_callback = TensorBoardFix(log_dir=log_dir, histogram_freq=1, write_graph=True, update_freq=1)
model.fit(dataset, epochs=200, callbacks=[tensorboard_callback])

这解决了我的问题，现在我可以通过调用 tf.summary.experimental.get_step().[=13 进入我的模型=]

（此问题可能会在更高版本的 TensorFlow 中修复）

TensorBoard 中 tf.summary.* 操作的步数始终为 0

Steps of tf.summary.* operations in TensorBoard are always 0

python

deep-learning

keras

tensorboard

tensorflow2.0