[Python] Is there a way to generate the computation graph of a PyTorch model from just a...

Imagine a network like this:

import torch
import torch.nn as nn

class CoolCNN(nn.Module):
def __init__(self):
super(CoolCNN, self).__init__()
self.initial_conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
self.parallel_conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)
self.secondary_conv = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
self.max_pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.fully_connected1 = nn.Linear(32 * 8 * 8, 128)
self.output_layer = nn.Linear(128, 10)

def forward(self, x):
main_path = self.max_pool(torch.relu(self.initial_conv(x)))
parallel_path = self.max_pool(torch.relu(self.parallel_conv(x)))
x = (main_path + parallel_path) / 2
x = self.max_pool(torch.relu(self.secondary_conv(x)))
x = x.view(-1, 32 * 8 * 8)
x = torch.relu(self.fully_connected1(x))
x = self.output_layer(x)
return x


It has this tree-like structure:

CoolCNN
└── Forward Pass
├── parallel_path
│ ├── parallel_conv (Conv2d)
│ ├── ReLU Activation
│ └── max_pool (MaxPool2d)
│
├── main_path
│ ├── initial_conv (Conv2d)
│ ├── ReLU Activation
│ └── max_pool (MaxPool2d)
│
├── Average main_path and parallel_path
│
├── secondary_conv (Conv2d)
├── ReLU Activation
└── max_pool (MaxPool2d)
│
├── Flatten the Tensor
│
├── fully_connected1 (Linear)
├── ReLU Activation
│
└── output_layer (Linear)


As you can see, the main and parallel paths being parallel branches gets accurately represented in the above illustrated tree representation.

I wanted to know if it's possible to somehow extract this tree-like structure from just a forward pass through the network without relying on the backward pass at all. There are libraries like torchviz which use the computation graph from backward pass which I don't want to use.

The only approach that I could find was forward hooks. That is in fact how torchsummary prints the model's summary.

However, with forward hooks, all you get to know is which order the nodes get called in. So it is just a topological sort of the underlying graph and it's impossible to accurately reconstruct a graph from just its topological sort. This means it's impossible to accurately derive the actual unique tree representation from forward hooks alone (at least that is my understanding).

I have already seen this question where the OP seemed to be happy without an exact graph reconstruction, so it doesn't address my need.

So is there a way by which we can reconstruct the computation graph from a forward pass alone?

Continue reading...