torchsummary

library1.5.1pypypiunverified

`torchsummary` provides a Keras-like `model.summary()` functionality for PyTorch models, displaying layer names, output shapes, parameter counts, and trainable parameters. It helps in quickly understanding the architecture and memory footprint of a neural network. The current version is 1.5.1, with releases occurring as needed for bug fixes and minor enhancements rather than a strict schedule.

ai-ml

pip install torchsummary

github ↗package ↗

INSTALL

IMPORT

SIG · TORCHSUMMARY

torchsummary

ai-mlpythonv1.5.1

Install

1.7s avg

Import

—

Disk

16MB

Pass rate

10/ 10

Env Coverage10 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v1.5.1 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.910 runs

installs and imports cleanly · install 0.0s · import 0.000s · 17.8MB

glibc

py 3.10–3.910 runs

installs and imports cleanly · install 1.7s · import 0.000s · 18MB

16MB installed

● package 16MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

summary

✓ from torchsummary import summary

✗ from torchsummary import summary

This example defines a simple Convolutional Neural Network and uses `torchsummary.summary` to print its architecture, output shapes, and parameter counts. Note the importance of providing a correct `input_size` tuple (excluding batch dimension) and ensuring the model is on the specified device.

import torch import torch.nn as nn from torchsummary import summary # Define a simple model 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(kernel_size=2) self.flatten = nn.Flatten() # Calculate input size for linear layer: (28 - 5 + 1) / 2 = 12 # So, 10 channels * 12 * 12 pixels = 1440 self.fc1 = nn.Linear(10 * 12 * 12, 50) self.relu2 = nn.ReLU() self.fc2 = nn.Linear(50, 10) def forward(self, x): x = self.pool1(self.relu1(self.conv1(x))) x = self.flatten(x) x = self.relu2(self.fc1(x)) x = self.fc2(x) return x model = SimpleCNN() # Determine device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # Print summary for an input image of 1 channel, 28x28 pixels # input_size should be (channels, height, width) *without* the batch dimension print(summary(model, input_size=(1, 28, 28), device=str(device)))

Debug

Known issues

gotchaThe `input_size` parameter must be a tuple representing the shape of a *single input sample* (excluding the batch dimension). For example, a batch of 64 images of size (3, 224, 224) would require `input_size=(3, 224, 224)`. Mismatches are a common cause of `RuntimeError` or incorrect summaries.

fix

Carefully determine the expected input shape for your model's first layer, excluding the batch dimension, and pass it as a tuple to `input_size`.

affects: All versions

gotcha`torchsummary` internally uses `torch.jit.trace` to analyze the model. Models with dynamic control flow (e.g., if-statements or loops whose behavior depends on input data values) may not be correctly summarized or might raise `RuntimeError` during tracing.

fix

For models with highly dynamic control flow, consider using alternative tools like `torchinfo` (which uses a different tracing mechanism) or manually inspecting layer outputs.

affects: All versions

gotchaThe `device` parameter in `summary()` expects a string (`"cpu"` or `"cuda"`) not a `torch.device` object directly.

fix

Convert your `torch.device` object to a string before passing it: `device=str(your_torch_device_object)`.

affects: All versions

Upgrade

Version history

1.5.1latest on PyPI · released Sep 26, 2018

Audit

Dependencies

tqdmrequiredUsed for progress bars internally, listed as a direct dependency.

torchrequiredPeer dependency, required for model definition and operations but not listed in install_requires as users typically install it separately.

Agent activity

17 hits · last 30 days

node

Amazon

OpenAI (training)

Resources

githubhttps://github.com/sksq96/pytorch-summary ↗

packagehttps://pypi.org/project/torchsummary/ ↗