Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages

, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

Repository files navigation

 ______ __ __ / __/ /___ ____ / /_/ /_ / /_/ / __ \/ __ \/ __/ __ \
/ __/ / /_/ / /_/ / /_/ / / /
/_/ /_/\____/ .___/\__/_/ /_/ /_/ 

flopth

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with cli tool and Python API.

Features

  • Handy cli command to show flops and params quickly
  • Visualization percent of flops and params in each layer
  • Support multiple inputs in model's forward function
  • Support Both CPU and GPU mode
  • Support Torchscript Model (Only Parameters are shown)
  • Support Python3.5 and above

Installation

Install stable version of flopth via pypi:

pip install flopth 

or install latest version via github:

pip install -U git+https://github.com/vra/flopth.git

Usage examples

cli command

flopth provide cli command flopth after installation. You can use it to get information of pytorch models quickly

Running on models in torchvision.models

with flopth -m <model_name>, flopth gives you all information about the <model_name>, input shape, output shape, parameter and flops of each layer, and total flops and params.

Here is an example running on alexnet (default input size in (3, 224, 224)):

$ flopth -m alexnet +---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===================+=============+=============+==========+==================+================================+==========+=================+=====================+
| features.0 | Conv2d | (3,224,224) | (64,55,55) | 23.296K | 0.0381271% | | 70.4704M | 9.84839% | #### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.1 | ReLU | (64,55,55) | (64,55,55) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.2 | MaxPool2d | (64,55,55) | (64,27,27) | 0.0 | 0.0% | | 193.6K | 0.027056% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.3 | Conv2d | (64,27,27) | (192,27,27) | 307.392K | 0.50309% | | 224.089M | 31.3169% | ############### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.4 | ReLU | (192,27,27) | (192,27,27) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.5 | MaxPool2d | (192,27,27) | (192,13,13) | 0.0 | 0.0% | | 139.968K | 0.0195608% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.6 | Conv2d | (192,13,13) | (384,13,13) | 663.936K | 1.08662% | | 112.205M | 15.6809% | ####### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.7 | ReLU | (384,13,13) | (384,13,13) | 0.0 | 0.0% | | 64.896K | 0.00906935% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.8 | Conv2d | (384,13,13) | (256,13,13) | 884.992K | 1.44841% | | 149.564M | 20.9018% | ########## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.9 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.10 | Conv2d | (256,13,13) | (256,13,13) | 590.08K | 0.965748% | | 99.7235M | 13.9366% | ###### |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.11 | ReLU | (256,13,13) | (256,13,13) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| features.12 | MaxPool2d | (256,13,13) | (256,6,6) | 0.0 | 0.0% | | 43.264K | 0.00604624% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| avgpool | AdaptiveAvgPool2d | (256,6,6) | (256,6,6) | 0.0 | 0.0% | | 9.216K | 0.00128796% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.0 | Dropout | (9216) | (9216) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.1 | Linear | (9216) | (4096) | 37.7528M | 61.7877% | ############################## | 37.7487M | 5.27547% | ## |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.2 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.3 | Dropout | (4096) | (4096) | 0.0 | 0.0% | | 0.0 | 0.0% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.4 | Linear | (4096) | (4096) | 16.7813M | 27.4649% | ############# | 16.7772M | 2.34465% | # |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.5 | ReLU | (4096) | (4096) | 0.0 | 0.0% | | 4.096K | 0.000572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
| classifier.6 | Linear | (4096) | (1000) | 4.097M | 6.70531% | ### | 4.096M | 0.572425% | |
+---------------+-------------------+-------------+-------------+----------+------------------+--------------------------------+----------+-----------------+---------------------+
FLOPs: 715.553M
Params: 61.1008M

Running on custom models

Also, given model name and the file path where the model defined, flopth will output model information:

For the dummpy network MyModel defined in /tmp/my_model.py,

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 to print model information:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 16.8591M
Params: 336.0

Multiple inputs

If your model has more than one input in forward, you can add multiple -i parameters to flopth:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

You can use flopth -m MyModel -p /tmp/my_model -i 3 224 224 -i 3 128 128 to print model information:

 flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Extra arguments in model's initialization

flopth with options like -x param1=int:3 param2=float:5.2 to process the extra parameters in model's initialization:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
# Please Notice the parameters ks1 and ks2 here!def__init__(self, ks1, ks2):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=ks1, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=ks2, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)

In order to pass value to the arguments of ks1 and ks2, we can run flopth like this:

$ flopth -m MyModel -p /tmp/my_model.py -i 3 224 224 -i 3 128 128 -x ks1=int:3 ks2=int:1
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+=======================+==========+=================+=========================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 43.75% | ##################### | 4.21478M | 47.6707% | ####################### |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv3 | Conv2d | (3,128,128) | (3,130,130) | 12 | 6.25% | ### | 202.8K | 2.29374% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
| conv4 | Conv2d | (3,130,130) | (3,132,132) | 12 | 6.25% | ### | 209.088K | 2.36486% | # |
+---------------+---------------+-------------+-------------+----------+------------------+-----------------------+----------+-----------------+-------------------------+
FLOPs: 8.84146M
Params: 192.0

Line number mode

One of the fancy features of flopth is that given the line number where the model object is definited, flopth can print model information:

# file path: /tmp/my_model.py# model name: MyModelimporttorch.nnasnnclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1, x2):
x1=self.conv1(x1)
x1=self.conv2(x1)
x2=self.conv3(x2)
x2=self.conv4(x2)
return (x1, x2)
if__name__=='__main__':
my_model=MyModel()

Since the model object my_model in defined in line 23, we can run flopth like this:

$ flopth -n 23 -p /tmp/my_model.py -i 3 224 224 -i 3 128 128
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 37.6923% | ################## |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,128,128) | (3,128,128) | 84 | 25.0% | ############ | 1.37626M | 12.3077% | ###### |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
FLOPs: 11.1821M
Params: 336.0

Notice: Although line number mode of flopth is quite handy, it may fail when the model definition is too complex, e.g., using outer config file to initialize a model. In this case, I recommend you to use flopth's Python API detailed below.

Python API

The Python API of flopth is quite simple:

importtorchimporttorch.nnasnn# import flopthfromflopthimportflopth# define ModelclassMyModel(nn.Module):
def__init__(self):
super(MyModel, self).__init__()
self.conv1=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv2=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv3=nn.Conv2d(3, 3, kernel_size=3, padding=1)
self.conv4=nn.Conv2d(3, 3, kernel_size=3, padding=1)
defforward(self, x1):
x1=self.conv1(x1)
x1=self.conv2(x1)
x1=self.conv3(x1)
x1=self.conv4(x1)
returnx1# declare Model objectmy_model=MyModel()
# Use input sizeflops, params=flopth(my_model, in_size=((3, 224, 224),))
print(flops, params)
# Or use input tensorsdummy_inputs=torch.rand(1, 3, 224, 224)
flops, params=flopth(my_model, inputs=(dummy_inputs,))
print(flops, params)

The output is like this:

16.8591M 336.0

To show detail information of each layer, add show_detail=True in flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), show_detail=True)

The outputs:

+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| module_name | module_type | in_shape | out_shape | params | params_percent | params_percent_vis | flops | flops_percent | flops_percent_vis |
+===============+===============+=============+=============+==========+==================+======================+==========+=================+=====================+
| conv1 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv2 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv3 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
| conv4 | Conv2d | (3,224,224) | (3,224,224) | 84 | 25.0% | ############ | 4.21478M | 25.0% | ############ |
+---------------+---------------+-------------+-------------+----------+------------------+----------------------+----------+-----------------+---------------------+
16.8591M 336.0

To show only the value of flops and params (no unit conversion), add bare_number=True to flopth function call:

flops, params=flopth(my_model, in_size=((3, 224, 224),), bare_number=True)

The outputs:

16859136 336

Known issues

  1. When use a module more than one time during forward, the FLOPs calculation is not correct, For example:
importtorch.nnasnnclassNet(nn.Module):
def__init__(self):
super(Net, self).__init__()
self.l1=nn.Linear(10, 10)
defforward(self, x, y):
x=self.l1(x)
x=self.l1(x)
x=self.l1(x)
returnx

Will give wrong FLOPs value, because we use register_buffer , which is bind to a nn.Module (in this example, l1).

TODOs

  • Support multiple inputs
  • Add parameter size
  • Add file line mode
  • Add line number mode
  • Support more modules

Contribution and issue

Any discussion and contribution are very welcomed. Please open an issue to reach me.

Acknowledge

This program is mostly inspired by torchstat, great thanks to the creators of it.

About

A simple program to calculate and visualize the FLOPs and Parameters of Pytorch models, with handy CLI and easy-to-use Python API.

Topics

Resources

Stars

131 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages