Hey, I’d like to know how can I draw/extract a model graph. I know that computation graph can give me some graph-like results, but each nodes can only be some tensor output by a specific differentible operation
In my case, I want the node of the graph to be the layer of the model instead of a tensor. Which means I just want to know the flow of the tensor from some layer to some other layer.
for example for a model defined like below:
class MyLayer(torch.nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.linear = torch.nn.Linear(in_dim, out_dim) self.relu = torch.nn.ReLU() def forward(x): return self.relu(self.linear(x)) class MyModel(torch.nn.Module): def __init__(self): super().__init__() self.my_layer1 = MyLayer(10, 5) self.my_layer2 = MyLayer(10, 5) self.my_layer3 = MyLayer(5, 3) def forward(self, x): x1 = self.my_layer1(x) x2 = self.my_layer2(x) x3 = x1 + x2 return self.my_layer3(x3)
I want a graph that shows the how the input is processed between “MyLayer” instance like below:
Is there any library or method I can achieve this?