如何在TensorFlow 2中使用tf.summary可视化每个步骤中tf.keras.layers.Layer或tf.keras.Model内部的图像张量？

如何解决如何在TensorFlow 2中使用tf.summary可视化每个步骤中tf.keras.layers.Layer或tf.keras.Model内部的图像张量？

在使用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()函数的参数 回调 中传递。

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

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

与@ tf.functions一起使用时，将在跟踪功能时捕获步骤值，因此除非使用tf.Variable步骤，否则对功能外部步骤的更改将不会反映在功能内部。 / p>

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

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

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

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

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

解决方法

我发现此问题已在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 方法，如下所示：

import { S3Client,GetObjectCommand,S3ClientConfig } from '@aws-sdk/client-s3';
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';


const s3Configuration: S3ClientConfig = {
    credentials: {
        accessKeyId: '<ACCESS_KEY_ID>',secretAccessKey: '<SECRET_ACCESS_KEY>'
    },region: '<REGION>',};
const s3 = new S3Client(s3Configuration);
const url = await getSignedUrl(s3,new GetObjectCommand({Bucket: '<BUCKET>',Key: '<KEY>' }),{ expiresIn: 15 * 60 }); // expires in seconds
console.log('Presigned URL: ',url);

并在 model.fit（）中使用此固定的回调类：

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,self).on_test_begin(*args,**kwargs)
    tf.summary.experimental.set_step(self._val_step)

这解决了我的问题，现在我可以通过调用 tf.summary.experimental.get_step（）进入模型内部。

（此问题可能在更高版本的TensorFlow中得以解决）