get_graph_node_names

torchvision.models.feature_extraction.get_graph_node_names(model: Module, tracer_kwargs: Optional[Dict[str, Any]] = None, suppress_diff_warning: bool = False) → Tuple[List[str], List[str]]

Dev utility to return node names in order of execution. See note on node names under create_feature_extractor(). Useful for seeing which node names are available for feature extraction. There are two reasons that node names can’t easily be read directly from the code for a model:

1. Not all submodules are traced through. Modules from torch.nn all fall within this category.

2. Nodes representing the repeated application of the same operation or leaf module get a _{counter} postfix.

The model is traced twice: once in train mode, and once in eval mode. Both sets of node names are returned.

For more details on the node naming conventions used here, please see the relevant subheading in the documentation.

Parameters:

• model (nn.Module) – model for which we’d like to print node names

• tracer_kwargs (dict, optional) – a dictionary of keyword arguments for NodePathTracer (they are eventually passed onto torch.fx.Tracer). By default it will be set to wrap and make leaf nodes all torchvision ops: {“autowrap_modules”: (math, torchvision.ops,),”leaf_modules”: _get_leaf_modules_for_ops(),} WARNING: In case the user provides tracer_kwargs, above default arguments will be appended to the user provided dictionary.

• suppress_diff_warning (bool, optional) – whether to suppress a warning when there are discrepancies between the train and eval version of the graph. Defaults to False.

Returns:

a list of node names from tracing the model in train mode, and another from tracing the model in eval mode.

Return type:

tuple(list, list)

Examples:

>>> model = torchvision.models.resnet18()
>>> train_nodes, eval_nodes = get_graph_node_names(model)