在 tensorflow-2.0 beta 中绘制子类模型答案

【问题标题】：Plotting subclassed model in tensorflow-2.0 beta在 tensorflow-2.0 beta 中绘制子类模型
【发布时间】：2019-12-16 16:56:54
【问题描述】：

我有一个子类模型，它通过子类化实例化一些自定义层。我尝试使用keras.utils.plot_model()，但它所做的只是打印模型块，没有出现任何图层。

Tensorflow 专家可以对此发表评论吗？将来会实现此功能吗？如果不是，那么检查计算图的下一个最佳选择是什么？注意model.summary()只给出了自定义层的参数汇总，其中包含两个密集层。理想情况下，我喜欢查看所有计算，如果这不是要求太多的话......

更新：我挖掘了源代码，看起来像 plot_model() 首先检查 _is_graph_network 属性。图网络用于功能和顺序 API。来源：

存在两种类型的Networks：图网络和子类网络。图形网络用于 Keras 功能和顺序 API。子类化当用户子类化Model 类时使用网络。一般来说， Graph Networks 支持的 Keras 功能比 Subclassed 多网络，特别是：

模型克隆 (keras.models.clone())

序列化（model.get_config()/from_config()，model.to_json()/to_yaml()）

全模型保存 (model.save())

(自定义图形组件) 自然，我想知道我是否可以构建图网络组件，以便我的子类模型/层可以使用这些功能。这需要很多努力吗？

(tf.function 图形可视化) 有人可以告诉我通过 Tensorboard 进行的图形可视化是否适用于 Tensorflow2 tf.functions？在 Tensorflow 1.x 中，定义了一组逻辑操作的名称范围（例如 GAN 中的生成器/鉴别器、VAE 中的编码器/解码器和损失/度量），然后它们在图形可视化中显示为高级块.我可以为 tf.functions 定义类似的东西吗？

【问题讨论】：

如果有人看到这个并且认识可以提供帮助的人，请告诉我。我真的希望能更多地了解 tensorflow2 的内部工作原理。谢谢。
你知道如何绘制使用子类化技术创建的模型吗？
@DavidKWH 如果你还在寻找答案，我已经发布了一个简单的解决方法here。

标签： python tensorboard tensorflow2.0

【解决方案1】：

根据官方文档https://www.tensorflow.org/tensorboard/graphs，你可以

使用 TensorFlow Summary Trace API 记录签名的函数 TensorBoard 中的可视化

这是一个可视化子类模型的简单示例：

import tensorflow as tf

from tensorflow.keras import layers

class MyModel(tf.keras.Model):

  def __init__(self, num_classes=10):
    super(MyModel, self).__init__(name='my_model')
    self.num_classes = num_classes
    self.dense_1 = layers.Dense(32, activation='relu')
    self.dense_2 = layers.Dense(num_classes, activation='sigmoid')

  def call(self, inputs): 
    x = self.dense_1(inputs)
    return self.dense_2(x)

model = MyModel(num_classes=10)

model.compile(optimizer=tf.keras.optimizers.RMSprop(0.001),
              loss='categorical_crossentropy',
              metrics=['accuracy'])

data = np.random.random((1000, 32))
labels = np.random.random((1000, 10))

@tf.function
def trace():
  data = np.random.random((1, 32))
  model(data)


logdir = "trace_log"
writer = tf.summary.create_file_writer(logdir)
tf.summary.trace_on(graph=True, profiler=True)
# Forward pass
trace()
with writer.as_default():
  tf.summary.trace_export(name="model_trace", step=0, profiler_outdir=logdir)

然后你可以使用 Tensorboard 来检查计算图：

tensorboard --logdir trace_log

【讨论】：