Tensorflow 2.4中如何可视化自定义模型的图形

Question

我希望能够在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：必须在导出之前启用跟踪。 我无法确定在哪里添加tf.summary.trace_on(graph=True)命令由TF提到。
有什么建议吗？

- Crysers

---

你找到解决方案了吗？我也有同样的问题。我正在尝试连接TensorBoard以进行一些配置文件，但是很困难，因为涉及到自定义类。 - StarShine

很遗憾，没有。但由于时间限制，我没有在这个问题上花费更多的精力。然而，如果您（或其他任何人）找到了解决方案，我将非常感兴趣用于未来的项目。 - Crysers

Answer 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。

- peachthief