在 Tensorflow 2.4 中可视化自定义模型的图形答案

【问题标题】：Visualize the graph of a custom model in Tensorflow 2.4在 Tensorflow 2.4 中可视化自定义模型的图形
【发布时间】：2021-04-18 17:53:45
【问题描述】：

我想在 Tensorboard 中可视化我的模型图，以检查我实现的模型是否正确。
我通过继承tf.keras.Layer 和tf.keras.Model 类来实现我自己的模型，看起来像这样（保留所有不必要的东西）：

class My_Model(tf.keras.Model):
    def __init__(self):
        super(Model_C_1, self).__init__()
        # actually here is much more, but this is not important.

    def build(self, inputs_shape):
        self.conv1 = tf.keras.layers.Conv1D(filters=16)
        # actually here is much more, but this is not important. 

    @tf.function
    def call(self, input, training):
        x = self.conv1(input)
        # actually here is much more, but this is not important. 
        return x

我想可视化计算图，因为模型要复杂得多，而且我不确定是否遗漏了什么（模型训练和工作，但我想仔细检查）。我的训练循环（非常简单）如下所示：

def train_step(batch, model, params, writer, optimizer):

    data = batch['data']
    with tf.GradientTape() as tape:
        predictions = model(data, training=True)
        loss = loss_object(labels, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    tf.summary.scalar(name='loss', data=loss, step=optimizer.iterations) 


# I process my data with pandas/numpy first and then
kf = KFold(n_splits=params.n_k_fold_splits)

for split_nr, (train_index, val_index) in enumerate(kf.split(ds)):
    #...
    writer = tf.summary.create_file_writer(params.path_train_log)
    for batch_train in train_ds:
        train_step(batch_train, model, params, writer, optimizer)

Tensorboard 中的训练和指标可视化效果很好。但是，我没有在 Tensorboard 中得到我的模型图。我在完整培训结束时尝试使用以下内容。

 with writer.as_default():
            tf.summary.trace_export(
              name="My_Trace",
              step=0)

这会导致错误：ValueError: Must enable trace before export. 我不知道在哪里添加tf.summary.trace_on(graph=True) 命令mentioned by TF。
有什么建议吗？

【问题讨论】：

您找到解决方案了吗？我几乎有同样的问题。自定义类并尝试连接 tensorboard 进行一些分析，但这很困难。
很遗憾没有。但是由于时间限制，我没有在这个问题上花费更多的精力。但是，如果您（或其他任何人）找到解决方案，我会对未来的项目非常感兴趣。

标签： python tensorflow keras tensorboard

【解决方案1】：

我发现让它工作的主要问题是你必须确保在你构建图表之前你正在做跟踪。因此，它必须在训练示例中的第一批运行，即您第一次调用模型时。

在我的示例中，train_step 是我的@tf.function

for batch, (input_image, target) in tqdm(train_ds.enumerate(), total=num_steps):
  if batch == 0:
    tf.summary.trace_on(graph=True)
    train_step(input_image, target, epoch)
    with summary_writer.as_default():
      tf.summary.trace_export(name='train_func', step=0)
  else:
    train_step(input_image, target, epoch)

如果您在它被调用之前不这样做，它似乎可以工作，创建一个大的 tf.events 文件但实际上并没有这样做。

另一种选择是在首次调用模型之前调用tf.summary.graph。

【讨论】：