Get all layers of the model in PyTorch

In this tutorial, we will learn about how to get all layers of the deep learning model in PyTorch.We will learn about the different methods we can use to access each layer in the model built in PyTorch.

We can install PyTorch using the pip command as follows:

pip install torch

It will install the latest version of pytorch on our system.

Once we install the PyTorch, we can use it to design to build deep learning architecture based on our problem at hand.

In the following tutorial, we are assuming you are dealing with some data and you have preprocessed it and now you are ready to design model architecture.

Note: Based on your problem statement we can have different model architects.

Here let’s consider the following model architecture:

import torch
import torch.nn as nn 
from torch_geometric.nn.conv import EdgeConv

class TableGNN(torch.nn.Module):

    def __init__(self):
        super(TableGNN, self).__init__()

        self.edgeconv1 = EdgeConv(dense_submodel_1, aggr='mean')
        self.edgeconv2 = EdgeConv(dense_submodel_2, aggr='mean')

        self.linear1 = torch.nn.Linear(256, 128)
        self.linear2 = torch.nn.Linear(128, 64)
        self.linear3 = torch.nn.Linear(64, 32)
        self.linear4 = torch.nn.Linear(32, 8)
        self.linear5 = torch.nn.Linear(8, 1)

        self.sigmoid = torch.nn.Sigmoid()
        self.relu = torch.nn.ReLU()

        self.norm1 = torch.nn.BatchNorm1d(256)
        self.norm2 = torch.nn.BatchNorm1d(128)
        self.norm3 = torch.nn.BatchNorm1d(64)
        self.norm4 = torch.nn.BatchNorm1d(32)
        self.norm5 = torch.nn.BatchNorm1d(8)

    def forward(self, data):
        edges = data.edge_index
        a = self.edgeconv1(x=data.x, edge_index=edges)
        a = self.edgeconv2(x=a, edge_index=edges)

        node_pair = torch.cat((a[edges.T[:, 0]], a[edges.T[:, 1]]), 1)
        x = self.norm1(node_pair)
        x = self.relu(x)

        x = self.linear1(x)
        x = self.norm2(x)
        x = self.relu(x)

        x = self.linear2(x)
        x = self.norm3(x)
        x = self.relu(x)

        x = self.linear3(x)
        x = self.norm4(x)
        x = self.relu(x)

        x = self.linear4(x)
        x = self.norm5(x)
        x = self.relu(x)

        x = self.linear5(x)
        x = self.sigmoid(x)

        return x

# Instantiate the model
model = TableGNN()

In the above, we are using the torch geometric library to design the model. We are training the model using graph neural networks. You can have model architecture based on CNN (convolutional Neural Network), RNN(Recurrent Neural Network), GAN ( Generative Adversarial Network).etc. depending on the problem statement.

Once we design model architecture we can get all layers in the model using different methods.

why accessing all the layers of the model is important?

1. To understand model architecture well so that we could easily modify it if required.
2. There are many different layers in the deep learning model. For Instance, convolutional layers, dense layers, pooling layers, etc. to understand the structure of each layer getting layers is important.
3. We can access individual layers which can help perform tasks like feature extraction, and interpreting the model.
4. Sometimes model architecture can be complex and it isn’t easy to get information about the individual layers by just reading. We can make it easy by printing all the layers at a time.

Different methods to get layers in PyTorch are as follows

children() and named_children() method

children()  method gives us layers of the model, named_children() method gives us the name of the layer and the layer itself in the model. It helps in recognizing the layers more easily than the children() method.

You can see the difference below:

Using children() method

for layer in model.children():
    print(layer)

Output:

Using named_children() method

for name, layer in model.named_children():
    print(name, layer)

Output:

We can see in the named_children() method the name of the layer for instance layer1, and layer2 has been printed along with the layer itself.

Using modules() and named_modules() methods

These methods help to print all the layers in the model. It can traverse through nested modules i.e. modules inside the modules. The Use of modules and named_modules is similar to that of the children and named_children method.

Using module method

for layer in model.modules():
    print(layer)

Output:

Using named_modules() method

for name, layer in model.named_modules():
    print(name, layer)

Output: