How to print a model in order of the invoking in forward function

I want to invoke a pre-trained model like this:

for module in model.children():
    x = module(x)

but sadly function model.children() does not return in order of the invoking in forward function.

Examples:
when I define my network like this:

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5,padding=(2,2))
        self.conv2 = nn.Conv2d(6, 16, 5,padding=(2,2))
        self.bn2 = nn.BatchNorm2d(16)
        self.bn1 = nn.BatchNorm2d(6)
        self.pool = nn.MaxPool2d(2, 2)    
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 7 * 7, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.pool(F.relu(x))
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.pool(F.relu(x))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        x = torch.softmax(x,dim=1)
        return x
    

function list(net.children()) output:

[Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)), Conv2d(6, 16, kernel_size=(5, 5), 
stride=(1, 1), padding=(2, 2)), BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True,
 track_running_stats=True), BatchNorm2d(6, eps=1e-05, momentum=0.1, affine=True, 
track_running_stats=True), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, 
ceil_mode=False), Linear(in_features=784, out_features=120, bias=True), Linear(in_features=120, 
out_features=84, bias=True), Linear(in_features=84, out_features=10, bias=True), ReLU()]

not the order defined in function forward but function __init__.

After search in Google, a package torchsummary is a good idea for this question, but if I invoke F.relu in forward function,

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 6, 28, 28]             156
       BatchNorm2d-2            [-1, 6, 28, 28]              12
         MaxPool2d-3            [-1, 6, 14, 14]               0
            Conv2d-4           [-1, 16, 14, 14]           2,416
       BatchNorm2d-5           [-1, 16, 14, 14]              32
         MaxPool2d-6             [-1, 16, 7, 7]               0
            Linear-7                  [-1, 120]          94,200
            Linear-8                   [-1, 84]          10,164
            Linear-9                   [-1, 10]             850
================================================================

There is no ReLU in results.

Any API about this function?

3 Likes

you can try register_forward_hook and print layer stats in the hook function