|
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 PackageInstallation: To install torchsummary, use pip:
pip install torchsummary Example : Here’s how you can use torchsummary to print the summary of a PyTorch model:
Python
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 SummaryIf you prefer not to use external packages, you can create a custom function to print the model summary. Here’s a basic implementation:
Python
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 torchinfoTo 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:
Python
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.
ConclusionPyTorch 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.
How to Print the Model Summary in PyTorch – FAQsWhat should I do if I encounter a shape mismatch error?A shape mismatch error most commonly according to the input size and the dimension of the first layer in the model. To correct this, ensure that the input tensor shape specified by the summary function corresponds to the input size acceptable by the first layer of the model.
Why are my custom layers not appearing in the model summary?In the model summary, custom layers might not be displayed if they have not been appropriately listed as subcategories of nn. Module. Make sure all the layers and components of your model that you created extend nn.
Can I print a summary of complex models and custom layers?Yes, you have the ability to print a summary of complex models and also of custom layers. However, it may be necessary to change the behavior of the custom summary function to get specific behaviors necessary or define extra attributes.
|
.webp){kind=small}
