[FEATURE] feature extraction stages

[jawi289o]

I have a question about understanding the structure and functionality of backbone stages. I want to know how to identify which layers belong to each stage.
For example, MobileNetV3 has 5 feature extraction stages (or downsampling stages), and I’d like to understand what layers or operations are included in each.
Stage 1: What kind of layers (e.g., convolutions, activations) are included?
Stage 2, Stage 3, etc.: What layers or operations are present here?
I want to better understand the breakdown of these stages to analyze how features are processed and extracted.

[rwightman]

@jawi289o you can use FX to extract graph node names, split them based on feature info module names

def split_nodes(nodes, split_names):
    """Split nodes at exact match if found, otherwise at last prefix match."""
    result = []
    start = 0

    for split in split_names:
        # First try exact match
        exact_matches = [i for i, node in enumerate(nodes) if node == split]

        if exact_matches:
            # If exact match found, split at first exact match
            end = exact_matches[0] + 1
        else:
            # Otherwise, split at last prefix match
            prefix_matches = [i for i, node in enumerate(nodes) if node.startswith(split)]
            if not prefix_matches:
                raise ValueError(f"No nodes found matching '{split}'")
            end = max(prefix_matches) + 1

        # Add segment and update start
        if end > start:  # Only add non-empty segments
            result.append(nodes[start:end])
            start = end

    # Add remaining nodes if any
    if start < len(nodes):
        result.append(nodes[start:])

    return result

feature_model = timm.create_model(model_name, features_only=True)
feat_names = feature_model.feature_info.module_name()

model = timm.create_model(model_name)
nodes = timm.models.get_graph_node_names(model)

split_nodes(nodes, feat_names)