Print model layer from which input is passed

1

In below code, input is passed from layer “self.linear1” in forward pass. I want to print the layers from which input is passed though other layer like “self.linear2” is initialise. It should be print only “linear1”.

import torch.nn as nn
import torch.nn.functional as F
import torch
from torchsummary import summary

class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(TwoLayerNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H) 
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
        h_relu = F.relu(self.linear1(x))
        #y_pred = self.linear2(h_relu)
        return h_relu

N, D_in, H, D_out = 32, 100, 50, 10

x = torch.randn(N, D_in)  
model = TwoLayerNet(D_in, H, D_out)
print("model is ",model)
y_pred = model(x)
2

A quick workaround will be using a wrapper layer over the pytorch’s linear layer where you print if input is flowing throw it.

Otherwise, you can simply add print statements or logging.

3

I just need the summary like printing model but the only layers through which input pass.

4

Maybe printing the graph in e.g TensorBoard as given in this tutorial might be helpful.

1 Like
5

I have also same issue any suggestion for this so please reply. Thanks in advance.

6

Above solution will help.