Skip to content
Estimates the size of a PyTorch model in memory
Jupyter Notebook Python
Branch: master
Clone or download
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
LICENSE.txt init Nov 18, 2017
README.md limitation clarification in README Mar 12, 2018
demo.ipynb init Nov 18, 2017
pytorch_modelsize.py fix Nov 22, 2017

README.md

PyTorch Model Size Estimator

This tool estimates the size of a PyTorch model in memory for a given input size.
Estimating the size of a model in memory is useful when trying to determine an appropriate batch size, or when making architectural decisions.

Note (1): SizeEstimator is only valid for models where dimensionality changes are exclusively carried out by modules in model.modules().

For example, use of nn.Functional.max_pool2d in the forward() method of a model prevents SizeEstimator from functioning properly. There is no direct means to access dimensionality changes carried out by arbitrary functions in the forward() method, such that tracking the size of inputs and gradients to be stored is non-trivial for such models.

Note (2): The size estimates provided by this tool are theoretical estimates only, and the total memory used will vary depending on implementation details. PyTorch utilizes a few hundred MB of memory for CUDA initialization, and the use of cuDNN alters memory usage in a manner that is difficult to predict. See this discussion on the PyTorch Forums for more detail.

See this blog post for an explanation of the size estimation logic.

Usage

To use the size estimator, simply import the SizeEstimator class, then provide a model and an input size for estimation.

# Define a model
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np

class Model(nn.Module):

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

        self.conv0 = nn.Conv2d(1, 16, kernel_size=3, padding=5)
        self.conv1 = nn.Conv2d(16, 32, kernel_size=3)

    def forward(self, x):
        h = self.conv0(x)
        h = self.conv1(h)
        return h

model = Model()

# Estimate Size
from pytorch_modelsize import SizeEstimator

se = SizeEstimator(model, input_size=(16,1,256,256))
print(se.estimate_size())

# Returns
# (size in megabytes, size in bits)
# (408.2833251953125, 3424928768)

print(se.param_bits) # bits taken up by parameters
print(se.forward_backward_bits) # bits stored for forward and backward
print(se.input_bits) # bits for input

Development

This tool is a product of the Laboratory of Cell Geometry at the University of California, San Francisco.

You can’t perform that action at this time.