Hello,
I’d like to determine the exact layer execution order (forward() 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.
Example: dummy Fashion-MNIST model, with a twist - I have swapped the declaration order of layer1 and layer2, whereas the forward() function contains the correct execution order.
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')

In another script, after loading back the model, I would like to determine the forward() call chain programmatically somehow:
...
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)
> 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))])

Is there a way to obtain the right order of the forward() function call chain? Can this be done somehow programmatically?
Maybe @ptrblck, could you please give me some hints? I’d really appreciate your help!
Thank you & Best regards,
RB