如何从Keras自定义子类层提取所有层并构建DAG

我尝试从继承自tf.keras.layers.Layer的自定义层中获取所有内部层，但遇到了困难。最终目标是构建包含tf.keras.layers.*基础层及运算的有向无环图（DAG）。

示例代码 #

from tensorflow import keras

class ResidualBlock(keras.layers.Layer):
    def __init__(
        self,
        filters0: int,
        filters1: int,
        activation: str = "leaky_relu",
        **kwargs,
    ) -> None:
        super().__init__()
        self.conv0 = keras.layers.Conv2D(
            filters0, 1, activation=activation, **kwargs)
        self.conv1 = keras.layers.Conv2D(
            filters1, 3, activation=activation,  **kwargs)

    def call(self, inputs, training=False):
        x = self.conv0(inputs, training=training)
        x = self.conv1(x, training=training)
        x = inputs + x
        return x

rb = ResidualBlock(2, 3)

new_model = keras.Sequential([rb, keras.layers.Dense(200)])

convert_to_DAG(new_model)

期望的DAG结构示例 #

[{'type': 'ResidualBlock', 'children': ['conv2D_1']}, 
 {'type': 'conv2D_1', 'children': ['conv2D_2', 'residual']}, 
 {'type': 'conv2D_2', 'children': ['Dense_1']}, 
 ...
]

已尝试的解决方案及不足 #

• 曾找到针对继承自tf.Model模型的方案，但不适用于继承自tf.keras.layers.Layer的自定义层；
• 另一个基于节点拆分模型的方案，无法递归追踪到基础层和运算，不符合需求。

最终需要生成的DAG需包含所有基础层及运算，结构类似包含基础层与残差运算等节点的完整图结构。

内容的提问来源于stack exchange，提问作者Hamza Kamran