Calculating flops of a given pytorch model

username: Ja-Keoung_Koo
cooked: <p>I am trying to made a function of calculating flops and want to discuss about it. In many papers, I can see the flop numbers, but it is hard to see the details of computing them.</p>
<p>I have some questions:</p>
<ul>
<li>Is it normal to include flops of ReLU, Batch normalization, …?</li>
<li>It seems common to consider the spatial dimension. For example, when calculating <code>Conv2d</code> layer, I need to know the image size. What is the common size of the original image size of width and height?</li>
<li>I am making a code as follows, but I wonder that there is a more elegant way to compute it. (In my code, it is hard to catch the dynamic spatial dimensions. For example, in <code>ReLU</code>, we don’t know the previous state.  )</li>
</ul>
<pre><code class="lang-auto">import torchvision
import re

def get_num_gen(gen):
    return sum(1 for x in gen)

def flops_layer(layer):
    """
    Calculate the number of flops for given a string information of layer.
    We extract only resonable numbers and use them.
    
    Args:
        layer (str) : example
            Linear (512 -&gt; 1000)
            Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
            BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True)
    """
    #print(layer)
    idx_type_end = layer.find('(')
    type_name = layer[:idx_type_end]
    
    params = re.findall('[^a-z](\d+)', layer)    
    flops = 1
    
    if layer.find('Linear') &gt;= 0:
        C1 = int(params[0])
        C2 = int(params[1])
        flops = C1*C2
        
    elif layer.find('Conv2d') &gt;= 0:
        C1 = int(params[0])
        C2 = int(params[1])
        K1 = int(params[2])
        K2 = int(params[3])
        
        # image size
        H = 32
        W = 32
        flops = C1*C2*K1*K2*H*W
    
#     print(type_name, flops)
    return flops

def calculate_flops(gen):
    """
    Calculate the flops given a generator of pytorch model.
    It only compute the flops of forward pass.
    
    Example:
        &gt;&gt;&gt; net = torchvision.models.resnet18()
        &gt;&gt;&gt; calculate_flops(net.children())
    """
    flops = 0;
    
    for child in gen:
        num_children = get_num_gen(child.children())
        
        # leaf node
        if num_children == 0:
            flops += flops_layer(str(child))
        
        else:
            flops += calculate_flops(child.children())
    
    return flops

net = torchvision.models.resnet18()
flops = calculate_flops(net.children())
print(flops / 10**9, 'G')
# 11.435429919 G
</code></pre>

username: smth
cooked: <p>for resnets, the spatial dimension is 224 height and 224 width. For inceptionv3 it is 299x299.<br>
Generally, since majority of flops are in conv and linear, nflops ~= X might show that you are approximating it, and that is prob sufficient for almost all things.</p>

username: YossiB
cooked: <p>You are missing some critical issue in your implementation: the dimension of the conv layers input is changing through the model depth.<br>
So i guess the right way to compute number of flops will be by using forward hook.</p>

username: warmspringwinds
cooked: <p>If someone still needs this, we wrote a small script to do that:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.py" target="_blank" rel="nofollow noopener">warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.py</a></h4>
<pre><code class="lang-py">import torch


# ---- Public functions

def add_flops_counting_methods(net_main_module):
    """Adds flops counting functions to an existing model. After that
    the flops count should be activated and the model should be run on an input
    image.
    
    Example:
    
    fcn = add_flops_counting_methods(fcn)
    fcn = fcn.cuda().train()
    fcn.start_flops_count()
    

    _ = fcn(batch)
    
    fcn.compute_average_flops_cost() / 1e9 / 2 # Result in GFLOPs per image in batch
</code></pre>

  This file has been truncated. <a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Example of usage:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/d5df5e066fe9c6078d38b26527d93436bf869b1c/pytorch_segmentation_detection/recipes/pascal_voc/segmentation/flops_counter.ipynb" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/d5df5e066fe9c6078d38b26527d93436bf869b1c/pytorch_segmentation_detection/recipes/pascal_voc/segmentation/flops_counter.ipynb" target="_blank" rel="nofollow noopener">warmspringwinds/pytorch-segmentation-detection/blob/d5df5e066fe9c6078d38b26527d93436bf869b1c/pytorch_segmentation_detection/recipes/pascal_voc/segmentation/flops_counter.ipynb</a></h4>
<pre><code class="lang-ipynb">{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/daniil/repos/anaconda2/lib/python2.7/site-packages/torch/nn/modules/module.py:482: UserWarning: src is not broadcastable to dst, but they have the same number of elements.  Falling back to deprecated pointwise behavior.\n",
      "  own_state[name].copy_(param)\n"
     ]
    }
   ],
   "source": [
    "import sys, os\n",
    "sys.path.append(\"/home/daniil/repos/pytorch-segmentation-detection/\")\n",
    "sys.path.insert(0, '/home/daniil/repos/pytorch-segmentation-detection/vision/')\n",
</code></pre>

  This file has been truncated. <a href="https://github.com/warmspringwinds/pytorch-segmentation-detection/blob/d5df5e066fe9c6078d38b26527d93436bf869b1c/pytorch_segmentation_detection/recipes/pascal_voc/segmentation/flops_counter.ipynb" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


username: niaoyu
cooked: <p>I just wrote a simple script to calculate the flops,.</p>
<p><a href="https://zhuanlan.zhihu.com/p/33992733" class="onebox" target="_blank" rel="nofollow noopener">https://zhuanlan.zhihu.com/p/33992733</a><br>
You can search the function named “print_model_parm_flops”</p>

username: thedch
cooked: <p><code>model.apply(fn)</code> combined with <code>module.register_forward_hook(hook)</code> allows for easy tracking of layers, but only works for <code>nn.Module</code> layers (conv, batchnorm, etc). This is effective for the majority of cases, but does not allow for tracking of functional calls, e.g. <code>F.interpolate(...)</code>. Is there any way to detect functional calls in a forward pass?</p>

username: bnikolic
cooked: <p>Coming to this rather late, but in case people are interested:</p>
<p>It is possible to directly measure the floating point operation count of models directly using CPU performance monitoring units as an alternative to the approaches which track the FLOPS of each operation. Using the python-papi module this is quite easy to do and the results match the operation counting as implemented by the thop module: see <a href="http://www.bnikolic.co.uk/blog/python/flops/2019/10/01/pytorch-count-flops.html" rel="nofollow noopener">http://www.bnikolic.co.uk/blog/python/flops/2019/10/01/pytorch-count-flops.html</a> for a comparison</p>