Model summary in PyTorch

1.8.0 · active · verified Mon Apr 13

torchinfo is a Python library that provides a detailed summary of PyTorch models, similar to the functionality found in TensorFlow's `model.summary()`. It displays information such as layer types, input/output shapes, kernel sizes, number of parameters, and computational operations (Mult-Adds). The library is a modern, re-written successor to older projects like `torchsummary` and `torchsummaryX`. The current version is 1.8.0, released in May 2023. While the last release is from May 2023, the project is considered actively developed and maintained within the PyTorch ecosystem.

Warnings

breaking Python 3.6 support was deprecated starting from version 1.6.0. Users on Python 3.6 should install `torchinfo` v1.5.4 or an earlier version.
Fix: Upgrade to Python 3.7+ or pin `torchinfo` to version 1.5.4 or lower (`pip install torchinfo==1.5.4`).
gotcha When using `input_size` without explicitly specifying `dtypes`, `torchinfo` defaults to `torch.float`. This can lead to `RuntimeError` for models expecting integer tensor inputs (e.g., embedding layers in NLP models like HuggingFace Transformers).
Fix: Explicitly pass the correct `dtypes` parameter to `summary()`, e.g., `summary(model, input_size=(batch, seq_len), dtypes=[torch.long])` for models with embedding layers.
gotcha `torchinfo` attempts to infer the device for model and input data, but mismatches between the model's device and the `input_data`'s device can cause `RuntimeError` (e.g., 'Expected all tensors to be on the same device').
Fix: Ensure both the model and the `input_data` (or the inferred `input_size` for `summary`) are on the same `torch.device`. You can explicitly pass the `device` parameter to `summary()` or move your model/input tensors to the correct device using `.to(device)`.
gotcha In Jupyter Notebooks or Google Colab, the output of `summary(model, ...)` may not display correctly unless explicitly wrapped in a `print()` statement.
Fix: Call `print(summary(model, ...))` to ensure the model summary is rendered in the notebook output.
gotcha Enabling `cache_forward_pass=True` for performance can lead to incorrect or outdated summaries if the model architecture or input data/sizes are modified without invalidating the cache.
Fix: Disable `cache_forward_pass` (default `False`) or be mindful to restart the kernel/process if making architectural changes while caching is enabled.
gotcha For 1D input tensors, `input_size` must be specified as a tuple with a trailing comma, e.g., `(10,)` instead of `(10)`, to be correctly interpreted.
Fix: Always use `(value,)` for single-dimension input sizes.
gotcha Version 1.8.0 introduced the use of `tensor.untyped_storage()` for PyTorch 2.0. While intended for compatibility, ensure that your PyTorch version is adequately supported for features that rely on this.
Fix: It is recommended to use `torchinfo` v1.8.0+ with PyTorch 2.0+ for full compatibility with these specific optimizations.

Install

pip install torchinfo Pip

Imports

summary
```
from torchinfo import summary
```

Quickstart

This example defines a simple Convolutional Neural Network (CNN) using `torch.nn.Module` and then uses `torchinfo.summary` to print a summary of its architecture, including layer shapes and parameter counts. The `input_size` parameter is crucial for `torchinfo` to perform a forward pass and calculate intermediate shapes.

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

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        # Calculate input features for the linear layer
        # ((28 - 5 + 1) / 2 - 5 + 1) / 2 = ((24 / 2) - 4) / 2 = (12 - 4) / 2 = 8 / 2 = 4
        # Output channels for conv2 is 20, so 20 * 4 * 4 = 320
        self.fc = nn.Linear(320, 10)

    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))
        x = x.view(-1, 320) # Flatten
        x = self.fc(x)
        return x

model = SimpleCNN()
# Provide an example input_size (batch_size, channels, height, width)
summary(model, input_size=(1, 1, 28, 28))

view raw JSON →