How to obtain all layers within a custom layer subclassed from keras.layers

Question

I'm trying to get all layers from within a custom layer subclassed from tf.keras.layers but I am having difficulties with this. The end goal is to create a DAG (Directed Acyclic Graph) that has layers from tf.keras.layers.* Here is an example:

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 = Sequential([rb, keras.layers.Dense(200)])

convert_to_DAG(new_model)

I want to get something like this:

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

I've seen all related answers like: How to access recursive layers of custom layer in tensorflow keras which accesses layers from a model subclassed from tf Model, NOT tf layers.Layer

The following code from Check which are the next layers in a tensorflow keras model which breaks a model based on nodes but it does not recursively follow each layer to its base layer/operations (which I require).

def get_layer_summary_with_connections(layer, relevant_nodes): 
    info = {}
    connections = []
    for node in layer._inbound_nodes: 
        if relevant_nodes and node not in relevant_nodes:
            continue
        for inbound_layer, node_index, tensor_index, _ in node.iterate_inbound(): 
            connections.append(inbound_layer.name)
    name = layer.name
    info['type'] = layer.__class__.__name__
    info['parents'] = connections
    return info

The end result should be a DAG that contains all base layers + operations like this: DAG end result. All Layers are base layers + operations

Thank you for any help. I can clarify if anything is unclear

Do you need the ResidualBlock to be a layer class? I've looked at the code of ResNet and their block is a function which maps the residual connections. But they also use the functional API for that. — mhenning
– mhenning, Commented Sep 1, 2023 at 7:30
@mhenning yes, ideally inheriting from the layer class since many models have already been compiled using layers that subclass the layer class. In an ideal case, I could just expand layers recursively using model.layers but that is not possible in this case. — Hamza Kamran
– Hamza Kamran, Commented Sep 1, 2023 at 18:11

MKimiSH · Accepted Answer · 2023-09-01 12:17:38Z

1

Have a look at tf.Module.submodules, which should be a list of all tf.Modules that are attributes of the current module, and it searches recursively. Both Layer and Model inherit from Module.

In your case, try adding the following to list the submodules and confirm that the submodules of the ResidualBlock are indeed submodules of the model.

print(new_model.submodules)
print(rb.submodules)

for rbs in rb.submodules:
    print("module: ", rbs, "isInModel?", rbs in new_model.submodules)

With this information it should be possible to create a DAG.

answered Sep 1, 2023 at 12:17

MKimiSH

1311 silver badge4 bronze badges

Sign up to request clarification or add additional context in comments.

5 Comments

Hamza Kamran Over a year ago

Thank you, but in this case, models subclassing keras.layers.* have already been instantiated and compiled. I would like to recursively follow the layers from layers.* subclassing, rather than module subclassing.

MKimiSH Over a year ago

@HamzaKamran I'm not so sure about this instantiation issue, rb.submodules can list the submodules (which are typically layers) without having to be in a Model. On the other hand, after looking closely at the image in the question, is it possible that your goal is actually to connect not just the Layers but also the operations in the computation graph? Because residual is not a Layer but a simple Add op.

Hamza Kamran Over a year ago

yes, operations are another part along with the layers. This is difficult and perhaps requires another question. I might be able to get this from the computation graph but I'm not sure how to

Hamza Kamran Over a year ago

yes operations are also required

MKimiSH Over a year ago

@HamzaKamran In that case, I think we are getting closer to what tensorboard does. Perhaps this is another question about "how to parse TF computation graph"? I'm unfortunately not well-versed in that realm.

Collectives™ on Stack Overflow

How to obtain all layers within a custom layer subclassed from keras.layers

1 Answer 1

5 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

5 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related