ONNX Hierarchical Trace

rashedkoutayni: <p>I have a simple model that consists of two parts:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Part1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3,16,3)
        self.conv2 = nn.Conv2d(16,32,3)
        self.maxpool1 = nn.MaxPool2d(32,2)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.maxpool1(x)
        return x


class Part2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(32,64,3)
        self.conv2 = nn.Conv2d(64,128,3)
        self.maxpool1 = nn.MaxPool2d(64,2)
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.maxpool1(x)
        return x


class Wrapper(nn.Module):
    def __init__(self):
        super().__init__()
        self.part1 = Part1()
        self.part2 = Part2()
    def forward(self, x):
        x = self.part1(x)
        x = self.part2(x)
        return x
</code></pre>
<p>If I trace the Wrapper model using ONNX:</p>
<pre><code class="lang-auto">dummy_input = torch.randn(1,3,224,224)
model = Wrapper()
torch.onnx.export(model, dummy_input, "wrapper.onnx", verbose=True, opset_version=11)
</code></pre>
<p>I get the following output:</p>
<pre><code class="lang-auto">graph(%input.1 : Float(1:150528, 3:50176, 224:224, 224:1, requires_grad=0, device=cpu),
      %part1.conv1.weight : Float(16:27, 3:9, 3:3, 3:1, requires_grad=1, device=cpu),
      %part1.conv1.bias : Float(16:1, requires_grad=1, device=cpu),
      %part1.conv2.weight : Float(32:144, 16:9, 3:3, 3:1, requires_grad=1, device=cpu),
      %part1.conv2.bias : Float(32:1, requires_grad=1, device=cpu),
      %part2.conv1.weight : Float(64:288, 32:9, 3:3, 3:1, requires_grad=1, device=cpu),
      %part2.conv1.bias : Float(64:1, requires_grad=1, device=cpu),
      %part2.conv2.weight : Float(128:576, 64:9, 3:3, 3:1, requires_grad=1, device=cpu),
      %part2.conv2.bias : Float(128:1, requires_grad=1, device=cpu)):
  %9 : Float(1:788544, 16:49284, 222:222, 222:1, requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%input.1, %part1.conv1.weight, %part1.conv1.bias)
  %10 : Float(1:1548800, 32:48400, 220:220, 220:1, requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%9, %part1.conv2.weight, %part1.conv2.bias) 
  %11 : Float(1:288800, 32:9025, 95:95, 95:1, requires_grad=1, device=cpu) = onnx::MaxPool[ceil_mode=0, kernel_shape=[32, 32], pads=[0, 0, 0, 0], strides=[2, 2]](%10)
  %12 : Float(1:553536, 64:8649, 93:93, 93:1, requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%11, %part2.conv1.weight, %part2.conv1.bias)
  %13 : Float(1:1059968, 128:8281, 91:91, 91:1, requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%12, %part2.conv2.weight, %part2.conv2.bias)
  %14 : Float(1:25088, 128:196, 14:14, 14:1, requires_grad=1, device=cpu) = onnx::MaxPool[ceil_mode=0, kernel_shape=[64, 64], pads=[0, 0, 0, 0], strides=[2, 2]](%13)
  return (%14)
</code></pre>
<p>which is fine.</p>
<p>But as you can see, the model is originally hierarchical as it consists of parts. However, this hierarchy is somehow flattened in the trace and there is no way to tell the part to which a layer belongs.</p>
<p>I might be able to know to which part each conv layer belongs because the parameters indicate this (e.g. part1.conv1.weight) but this is not guaranteed and I can’t find out from the trace which part the first pooling layer belongs to.</p>
<p>Is there a way to keep the hierarchy while tracing or at least indicate it?</p>

rashedkoutayni: <p>In other words, is there a way to change %9, %10, etc to something that indicates to which part each layer belongs (is it even possible to change them anyway?)</p>

rashedkoutayni: <p>Looks like the scope is no more available since PyTorch 1.4<br>
There is a workaroud to get the scope mentioned here:</p><aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/33463" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/33463" target="_blank" rel="noopener nofollow ugc">torch._C.Node.scopeName() missing in pytorch 1.4</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-02-18" data-time="21:44:40" data-timezone="UTC">09:44PM - 18 Feb 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/xvdp" target="_blank" rel="noopener nofollow ugc">
          <img alt="xvdp" src="https://avatars.githubusercontent.com/u/40925114?v=4" class="onebox-avatar-inline" width="20" height="20">
          xvdp
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
scopeName() appears as a method in pytorch 1.4 torch._C.Node but it seems to lo longer be written to.
It is minor...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">oncall: jit</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
<p>
while iterating over the graph nodes, the scope name can be obtained using: node.scopeName()</p>