How to Print the Model Summary in PyTorch

Printing a model summary is a crucial step in understanding the architecture of a neural network. In frameworks like Keras, this is straightforward with the model.summary() method. However, in PyTorch, achieving a similar output requires a bit more work. This article will guide you through the process of printing a model summary in PyTorch, using the torchinfo package, which is a successor to torch-summary.

Why Model Summary is Important?

Before diving into the implementation, let's briefly discuss why having a model summary is important:

• Debugging: Helps in identifying issues with the model architecture.
• Optimization: Provides insights into the number of parameters and computational complexity.
• Documentation: Serves as a quick reference for the model architecture.
  

Step-by-Step Guide for Getting the Model Summary

'torchsummary' is a useful package to obtain the architectural summary of the model in the same similar as in case of Keras’ model. summary(). It shows the layer types, the resultant shape of the model, and the number of parameters available in the models.

1. Using torchsummary Package

Installation: To install torchsummary, use pip:

pip install torchsummary

Example : Here’s how you can use torchsummary to print the summary of a PyTorch model:

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

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(in_features=64*28*28, out_features=128)
        self.fc2 = nn.Linear(in_features=128, out_features=10)
        
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.view(x.size(0), -1)  # Flatten the tensor
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN()
summary(model, input_size=(1, 28, 28))

Output:

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 32, 28, 28]             320
            Conv2d-2           [-1, 64, 28, 28]          18,496
            Linear-3                  [-1, 128]       6,422,656
            Linear-4                   [-1, 10]           1,290
================================================================
Total params: 6,442,762
Trainable params: 6,442,762
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.58
Params size (MB): 24.58
Estimated Total Size (MB): 25.16
----------------------------------------------------------------

2. Custom Implementation for Model Summary

If you prefer not to use external packages, you can create a custom function to print the model summary. Here’s a basic implementation:

def print_model_summary(model, input_size):
    def register_hook(module):
        def hook(module, input, output):
            class_name = str(module.__class__).split(".")[-1].split("'")[0]
            module_idx = len(summary)
            m_key = f"{class_name}-{module_idx+1}"
            summary[m_key] = {
                "input_shape": list(input[0].size()),
                "output_shape": list(output.size()),
                "nb_params": sum(p.numel() for p in module.parameters())
            }
        if not isinstance(module, nn.Sequential) and not isinstance(module, nn.ModuleList) and module != model:
            hooks.append(module.register_forward_hook(hook))

    summary = {}
    hooks = []
    model.apply(register_hook)
    with torch.no_grad():
        model(torch.zeros(1, *input_size))

    for h in hooks:
        h.remove()

    print("----------------------------------------------------------------")
    line_new = "{:>20}  {:>25} {:>15}".format("Layer (type)", "Output Shape", "Param #")
    print(line_new)
    print("================================================================")
    total_params = 0
    for layer in summary:
        line_new = "{:>20}  {:>25} {:>15}".format(
            layer,
            str(summary[layer]["output_shape"]),
            "{0:,}".format(summary[layer]["nb_params"])
        )
        total_params += summary[layer]["nb_params"]
        print(line_new)
    print("================================================================")
    print(f"Total params: {total_params:,}")
    print("----------------------------------------------------------------")

# Example usage
print_model_summary(model, (1, 28, 28))

Output:

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [1, 32, 28, 28]             320
            Conv2d-2            [1, 64, 28, 28]          18,496
            Linear-3                   [1, 128]       6,422,656
            Linear-4                    [1, 10]           1,290
================================================================
Total params: 6,442,762
----------------------------------------------------------------

If you integrate the model to more complicated models or define more layers, you may need to make more changes in the custom summary of the function for the specific behaviors or some attributes that you added. Make certain that all submodules are correctly registered for the generation of the correct summary.

3. Using torchinfo

To print the model summary in PyTorch, we will use the torchinfo package. You can install it using pip:

pip install torchinfo

Basic Example of torchinfo:

import torch
import torch.nn as nn
from torchinfo import summary

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.fc1 = nn.Linear(64 * 5 * 5, 128)  # Updated in_features to 64 * 5 * 5
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # Apply pooling
        x = self.pool(torch.relu(self.conv2(x)))  # Apply pooling
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleModel()
summary(model, input_size=(1, 1, 28, 28))

Output:

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
SimpleModel                              [1, 10]                   --
├─Conv2d: 1-1                            [1, 32, 26, 26]           320
├─MaxPool2d: 1-2                         [1, 32, 13, 13]           --
├─Conv2d: 1-3                            [1, 64, 11, 11]           18,496
├─MaxPool2d: 1-4                         [1, 64, 5, 5]             --
├─Linear: 1-5                            [1, 128]                  204,928
├─Linear: 1-6                            [1, 10]                   1,290
==========================================================================================
Total params: 225,034
Trainable params: 225,034
Non-trainable params: 0
Total mult-adds (M): 2.66
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.24
Params size (MB): 0.90
Estimated Total Size (MB): 1.14
==========================================================================================

Common Issues in Model Summary Printing

• Shape Mismatch: A frequent mistake when printing the model summary is a shape mismatch. This mostly happens when the size of input data given does not meet the required dimension of the first layer of the model. To resolve this, check that the dimensionality of the input tensor is in accordance with the needed size in the first layer of the model. It is always recommended to verify the input size being passed to the summary function and the model’s first layer dimensions.
• Unregistered Modules: Another common scenario is unregistered modules, and more specifically, the custom layers or containers. These components must be classes inheriting from nn. Module must be correctly registered to the summary. If a custom component is not subclassing nn, This means that the documentation generated by derived classes to traverse and, if necessary, modify the defined architecture will contain incomplete information. Module, it will not be displayed in the summary of the model; Make sure every custom layer or module you use is defined as a subclass of nn. Module to do this and overcome this problem.

Conclusion

PyTorch is convenient in visualizing neural network architectures and debugging them through printing a model summary. Regardless of using the torchsummary or any other package or method of obtaining the model structure, clearly observing the model structure helps in development of the model and in troubleshooting.

By reading this tutorial, you should be able to install and import torchsummary successfully, and write a generally custom model summary function, and solve general problems and complex models.