Tracing a Graph of Torch Layers

Erotemic: <p>Given a torch module I would like to create a graph where the nodes correspond to torch modules (i.e. layer classes) and an edge between layer u and v means that the output of u eventually reaches the input of v (there might be intermediate functional operations between them).</p>
<p>Note, my definition of an edge isn’t exactly precise. What I really mean is I want to extract a directed graph who’s transitive closure has the afformentioned property. So in the graph I actually care about all I care about is that if the input of v depends on the output of u, there is a path between u and v.</p>
<p>First I attempted to use torch.jit.trace, but the problem was that I couldn’t figure out how to associate the nodes in the JIT trace to the named layers in the torch module. I had a similar problem when trying to use the backwards pass graph.</p>
<p>It seems that tools that visualize graphs (as discussed in this SO thread <a href="https://stackoverflow.com/questions/52468956/how-do-i-visualize-a-net-in-pytorch" class="inline-onebox" rel="noopener nofollow ugc">python - How do I visualize a net in Pytorch? - Stack Overflow</a>) like torchviz rely on these, but they also don’t seem to be able to grab the layer names they correspond to.</p>
<p>I attempted an alternative pass where I look at the inputs / outputs of layers using a forward hook and connecting up layers where the input had the same id as an output, but this misses connections if there is an intermediate functional operation between sublayer forward passes:</p>
<pre><code class="lang-auto">        from itertools import count
        from collections import defaultdict
        import torch

        def model_layers(model):
            """ Extract named "leaf" layers from a module """
            stack = [('', '', model)]
            while stack:
                prefix, basename, item = stack.pop()
                name = '.'.join([p for p in [prefix, basename] if p])
                if isinstance(item, torch.nn.modules.conv._ConvNd):
                    yield name, item
                elif isinstance(item, torch.nn.modules.batchnorm._BatchNorm):
                    yield name, item
                elif hasattr(item, 'reset_parameters'):
                    yield name, item

                child_prefix = name
                for child_basename, child_item in list(item.named_children())[::-1]:
                    stack.append((child_prefix, child_basename, child_item))

        # Create example network
        import torchvision
        net = torchvision.models.resnet18()

        # The layers are the nodes we care about
        named_layers = list(model_layers(net))

        inputs = torch.rand(2, 3, 224, 224)

        layer_io_ids = {}
        counter = count()

        # Use hooks to trace the inputs/outputs to extract some of the edges
        def _make_layer_hook(name):
            def record_layer_io(layer, args, output):
                iter_idx = next(counter)
                info = {
                    'iter_idx': iter_idx,
                }
                if len(args) == 1:
                    input_tensor = args[0]
                    info['input_ids'] = [id(input_tensor)]
                else:
                    raise NotImplementedError
                    input_tensor = None

                if isinstance(output, torch.Tensor):
                    info['output_ids'] = [id(output)]
                else:
                    raise NotImplementedError

                layer_io_ids[name] = info
            return record_layer_io

        for name, layer in named_layers:
            layer._forward_hooks.clear()
            hook = _make_layer_hook(name)
            layer.register_forward_hook(hook)

        # Forward pass to call the hooks
        outputs = net(inputs)

        input_id_to_name = defaultdict(list)
        output_id_to_name = defaultdict(list)
        for name, info in layer_io_ids.items():
            for input_id in info['input_ids']:
                input_id_to_name[input_id].append(name)
            for output_id in info['output_ids']:
                output_id_to_name[output_id].append(name)

        # Build the graph with the info we have
        import networkx as nx
        layer_graph = nx.DiGraph()

        for name, info in layer_io_ids.items():
            layer_graph.add_node(name)

        for name, info in layer_io_ids.items():
            output_ids = info['output_ids']
            for output_id in output_ids:
                if output_id in input_id_to_name:
                    for other in input_id_to_name[output_id]:
                        layer_graph.add_edge(name, other)

        nx.write_network_text(layer_graph, vertical_chains=1)
</code></pre>
<p>results in something close to what I want, but I need to figure out the connections between the current weakly connected components. My first idea was to look at the order in which forward was called (to get a topological sort of the graph I’m interested in), but it might be the case that a network has branching paths (e.g. multiple heads), so I can’t use that strategy.</p>
<pre><code class="lang-auto">╟── conv1
╎   ╽
╎   bn1
╟── layer1.0.conv1
╎   ╽
╎   layer1.0.bn1
╎   ╽
╎   layer1.0.conv2
╎   ╽
╎   layer1.0.bn2
╎   ╽
╎   layer1.1.conv1
╎   ╽
╎   layer1.1.bn1
╎   ╽
╎   layer1.1.conv2
╎   ╽
╎   layer1.1.bn2
╎   ├─╼ layer2.0.conv1
╎   │   ╽
╎   │   layer2.0.bn1
╎   │   ╽
╎   │   layer2.0.conv2
╎   │   ╽
╎   │   layer2.0.bn2
╎   │   ╽
╎   │   layer2.1.conv1
╎   │   ╽
╎   │   layer2.1.bn1 ╾ layer2.0.downsample.1
╎   │   ╽
╎   │   layer2.1.conv2
╎   │   ╽
╎   │   layer2.1.bn2
╎   │   ├─╼ layer3.0.conv1
╎   │   │   ╽
╎   │   │   layer3.0.bn1
╎   │   │   ╽
╎   │   │   layer3.0.conv2
╎   │   │   ╽
╎   │   │   layer3.0.bn2
╎   │   │   ╽
╎   │   │   layer3.1.conv1
╎   │   │   ╽
╎   │   │   layer3.1.bn1 ╾ layer3.0.downsample.1
╎   │   │   ╽
╎   │   │   layer3.1.conv2
╎   │   │   ╽
╎   │   │   layer3.1.bn2
╎   │   │   ├─╼ layer4.0.conv1
╎   │   │   │   ╽
╎   │   │   │   layer4.0.bn1
╎   │   │   │   ╽
╎   │   │   │   layer4.0.conv2
╎   │   │   │   ╽
╎   │   │   │   layer4.0.bn2
╎   │   │   │   ╽
╎   │   │   │   layer4.1.conv1
╎   │   │   │   ╽
╎   │   │   │   layer4.1.bn1 ╾ layer4.0.downsample.1
╎   │   │   │   ╽
╎   │   │   │   layer4.1.conv2
╎   │   │   │   ╽
╎   │   │   │   layer4.1.bn2
╎   │   │   └─╼ layer4.0.downsample.0
╎   │   │       ╽
╎   │   │       layer4.0.downsample.1
╎   │   │       └─╼  ...
╎   │   └─╼ layer3.0.downsample.0
╎   │       ╽
╎   │       layer3.0.downsample.1
╎   │       └─╼  ...
╎   └─╼ layer2.0.downsample.0
╎       ╽
╎       layer2.0.downsample.1
╎       └─╼  ...
╙── fc
</code></pre>
<p>All in all it seems like the JIT or backwards graph is the only way to really get the information flow  topology, but I’d really like to see it in terms of the original layer names and not the low-level functional components. Is there any way to associate the JIT or backwards graph back to layer names?</p>
<p>EDIT: it looks like torchview: <a href="https://github.com/mert-kurttutan/torchview" class="inline-onebox" rel="noopener nofollow ugc">GitHub - mert-kurttutan/torchview: torchview: visualize pytorch models</a> might do this? But maybe it doens’t get the layer names?</p>

vdw: <p>You could check out some of the blog posts by <a class="mention" href="/u/tom">@tom</a></p>
<ul>
<li><a href="https://lernapparat.de/interactive-pytorch-visualization" rel="noopener nofollow ugc">Interactive Visualization of PyTorch (JIT) models</a></li>
<li><a href="https://lernapparat.de/jit-python-graphops" rel="noopener nofollow ugc">Python graph operations in the JIT</a></li>
<li><a href="https://lernapparat.de/visualize-pytorch-models" rel="noopener nofollow ugc">Visualizing PyTorch model structure</a></li>
</ul>
<p>Those example use <code>torch.jit.trace</code> which is indeed not overly intuitive, but the blog posts do give good ideas.</p>
<p>You should also check out <a href="https://pytorch.org/docs/stable/fx.html" rel="noopener nofollow ugc"><code>torch.fx</code></a> which seem rather new and simpler / more intuitive than <code>torch.jit.trace</code>. But I haven’t really figured this one out either :).</p>

johnmarktaylor: <p>TorchLens does this, and remains in active development if there’s any other features on your wishlist:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/johnmarktaylor91/torchlens">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/johnmarktaylor91/torchlens" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:666/500;"><img src="https://repository-images.githubusercontent.com/551190643/a758b584-7a74-4bec-b83a-7bc2dab612e7" class="thumbnail" width="666" height="500"></div>

<h3><a href="https://github.com/johnmarktaylor91/torchlens" target="_blank" rel="noopener nofollow ugc">GitHub - johnmarktaylor91/torchlens: Package for extracting and mapping the...</a></h3>

  <p>Package for extracting and mapping the results of every single tensor operation in a PyTorch model in one line of code.  - GitHub - johnmarktaylor91/torchlens: Package for extracting and mapping th...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>