在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: Must enable trace before export. 我不知道在哪里添加tf.summary.trace_on(graph=True)命令mentioned by TF。
有什么建议吗

Answer 1

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

在我的例子中，train_step是my@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