Check which are the next layers in a tensorflow keras model

Viewed 553

I have a model which has shortcuts between layers. For each layer, I would like to get the name (or index) of the next connected layers, because simply iterating through all the model.layers will not tell me whether the layer was connected to the previous one or not.

An example model could be:

model = tf.keras.applications.resnet50.ResNet50(
include_top=True, weights='imagenet', input_tensor=None,
input_shape=None, pooling=None, classes=1000)
2 Answers

You can extract the information in dict format in this way...

Firstly, define a utility function and get the relevant nodes as made in the model.summary() method from every Functional model (code reference)

relevant_nodes = []
for v in model._nodes_by_depth.values():
    relevant_nodes += v

def get_layer_summary_with_connections(layer):
    
    info = {}
    connections = []
    for node in layer._inbound_nodes:
        if relevant_nodes and node not in relevant_nodes:
            # node is not part of the current network
            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

Secondly, extract the information iterating through layers:

results = {}
layers = model.layers
for layer in layers:
    info = get_layer_summary_with_connections(layer)
    results[layer.name] = info

results is a nested dict with this format:

{
  'layer_name': {'type':'the layer type', 'parents':'list of the parent layers'},
  ...
  'layer_name': {'type':'the layer type', 'parents':'list of the parent layers'}
}

For ResNet50 it results in:

{
  'input_4': {'type': 'InputLayer', 'parents': []},
  'conv1_pad': {'type': 'ZeroPadding2D', 'parents': ['input_4']},
  'conv1_conv': {'type': 'Conv2D', 'parents': ['conv1_pad']},
  'conv1_bn': {'type': 'BatchNormalization', 'parents': ['conv1_conv']},
  ...
  'conv5_block3_out': {'type': 'Activation', 'parents': ['conv5_block3_add']},
  'avg_pool': {'type': 'GlobalAveragePooling2D', 'parents' ['conv5_block3_out']},
  'predictions': {'type': 'Dense', 'parents': ['avg_pool']}
}

Also, you can modify get_layer_summary_with_connections to return all the information you are interested in

You can view the whole model and its connections with the keras's Model plotting utilities

tf.keras.utils.plot_model(model, to_file='path/to/image', show_shapes=True)
Related