时间戳

首页 > 解决方案 > TensorBoard 中 tf.summary.* 操作的步数始终为 0

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

问题描述

当我使用 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 的图像。

我尝试将参数step设置为从以下tf.summary.image()获得的值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()中传递。

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

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

将 this 与 @tf.functions 一起使用时，将在跟踪函数时捕获步长值，因此除非使用 tf.Variable 步长，否则对函数外步长的更改不会反映在函数内。

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

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

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

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

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

标签： pythonkerasdeep-learningtensorflow2.0tensorboard

解决方案

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

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

要修复它，请继承tf.keras.callbacks.TensorBoard类并覆盖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()中使用这个固定的回调类：

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()在我的模型中获得正确的步骤。

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

推荐阅读