Extracting hierarchy/structure information for operators with opcode of "call_function" or "call_method"

pimenton: <p>After tracing an <code>nn.Module</code> with torch.fx’s tracer, I can tell the hierarchical information of operator nodes with opcode of “call_module” by looking at their <code>target</code> attribute (the qualified name) of the operator node. For example, <code>SubMod.conv</code> tells that the convolution layer <code>conv</code> is within the <code>SubMod</code> module. However, I am not able to do something similar for operator nodes with opcode of “call_method”, “call_function”, “placeholder”, “output”, etc. The ones that I have particular interests in are the ones with opcode of “call_method” and “call_function”, and I am wondering if there’s any way to extract their location information in the hierarchical structure of a nested module. When I look at the <code>target</code> attribute of the operator node with opcode of “call_method” or “call_function”, I usually see <code>&lt;built-in function</code> and sometimes along with the hex address of the function or method object, but it’s hard to infer the structure information from that.</p>
<p>Being able to obtain the “call_method” and “call_function” operator nodes location information in the hierarchical structure is particularly essential and helpful for distinguishing models like the following two after tracing them with torch.fx’s tracer: the two models below (<code>Model1</code> and <code>Model2</code>) are structurally different in PyTorch code; yet after tracing, the printed graph/operator nodes tables for both of the modules are identical (please see screenshot below for the printed graph/operator nodes table output and also the code for reproducing this). I am wondering if there’s any good way to obtain this information. I would really appreciate any suggestions, thank you very much in advance!</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.fx as fx

class SubModule1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 3, (3, 3))
    
    def forward(self, x):
        x = self.conv(x)
        x = x + x
        return x

class Model1(nn.Module):
    def __init__(self):
        super().__init__()
        self.SubMod = SubModule1()

    def forward(self, x):
        x = self.SubMod(x)
        return x

class SubModule2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 3, (3, 3))

    
    def forward(self, x):
        x = self.conv(x)
        return x

class Model2(nn.Module):
    def __init__(self):
        super().__init__()
        self.SubMod = SubModule2()

    def forward(self, x):
        x = self.SubMod(x)
        x = x + x
        return x

def print_graph_table(module):
    tracer_class = fx.Tracer
    graph = tracer_class().trace(module)
    graph.print_tabular()

model1 = Model1()
print("model1 Graph Table:")
print_graph_table(model1)
print()

model2 = Model2()
print("model2 Graph Table:")
print_graph_table(model2)
</code></pre>
<h3><a name="printed-graphoperator-nodes-tables-output-for-both-model1-and-model2-from-running-the-code-above-1" class="anchor" href="#printed-graphoperator-nodes-tables-output-for-both-model1-and-model2-from-running-the-code-above-1"></a>Printed graph/operator nodes tables output for both <code>Model1</code> and <code>Model2</code> (from running the code above):</h3>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd5a619107b90b35ba60ea0e976a3968f2436858.png" data-download-href="https://discuss.pytorch.org/uploads/default/bd5a619107b90b35ba60ea0e976a3968f2436858" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/d/bd5a619107b90b35ba60ea0e976a3968f2436858.png" alt="image" data-base62-sha1="r15SIB5t5sG66S70LaMxEXpkfNK" width="690" height="317" data-dominant-color="252525"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">731×336 8.05 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>