Determine forward() call chain of RecursiveScriptModule

robeast: <p>Hello,<br>
I’d like to determine the exact layer execution order (<code>forward()</code> call chain) of a RecursiveScriptModule. I have to do it on this particular data type, as this is the pre-defined interface for model exchange in our flow.</p>
<p><strong>Example</strong>: dummy Fashion-MNIST model, with a twist - I have swapped the declaration order of <code>layer1</code> and <code>layer2</code>, whereas the <code>forward()</code> function contains the correct execution order.</p>
<pre><code class="lang-python">class FashionCNN(nn.Module):

    def __init__(self):
        super(FashionCNN, self).__init__()

        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.fc1 = nn.Linear(in_features=64*6*6, out_features=600)
        self.drop = nn.Dropout2d(0.25)
        self.fc2 = nn.Linear(in_features=600, out_features=120)
        self.fc3 = nn.Linear(in_features=120, out_features=10)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.view(out.size(0), -1)
        out = self.fc1(out)
        out = self.drop(out)
        out = self.fc2(out)
        out = self.fc3(out)
        return out

...
model = FashionCNN()
...
scripted_model = torch.jit.script(model.eval())
scripted_model.save('fmnist_scripted.pt')
</code></pre>
<p>In another script, after loading back the model, I would like to determine the <code>forward()</code> call chain programmatically somehow:</p>
<pre><code class="lang-python">...
scripted_model = torch.jit.load('fmnist_scripted.pt')

# Iterate over all hierarchycal layers ?!
layers = OrderedDict()
for i in scripted_model.named_modules():
    if not list(i[1].named_children()):
        layers[i[0]] = i[1]

# This mechanism gives me the wrong layer ordering:
print(layers)
&gt; OrderedDict([('layer2.0', RecursiveScriptModule(original_name=Conv2d)),
             ('layer2.1', RecursiveScriptModule(original_name=BatchNorm2d)),
             ('layer2.2', RecursiveScriptModule(original_name=ReLU)),
             ('layer2.3', RecursiveScriptModule(original_name=MaxPool2d)),
             ('layer1.0', RecursiveScriptModule(original_name=Conv2d)),
             ('layer1.1', RecursiveScriptModule(original_name=BatchNorm2d)),
             ('layer1.2', RecursiveScriptModule(original_name=ReLU)),
             ('layer1.3', RecursiveScriptModule(original_name=MaxPool2d)),
             ('fc1', RecursiveScriptModule(original_name=Linear)),
             ('drop', RecursiveScriptModule(original_name=Dropout2d)),
             ('fc2', RecursiveScriptModule(original_name=Linear)),
             ('fc3', RecursiveScriptModule(original_name=Linear))])
</code></pre>
<p>Is there a way to obtain the right order of the <code>forward()</code> function call chain? Can this be done somehow programmatically?</p>
<p>Maybe <a class="mention" href="/u/ptrblck">@ptrblck</a>, could you please give me some hints? I’d really appreciate your help!</p>
<p>Thank you &amp; Best regards,<br>
RB</p>