Calculating flops of a given pytorch model

warmspringwinds · March 9, 2018, 2:46am

If someone still needs this, we wrote a small script to do that:

warmspringwinds/pytorch-segmentation-detection/blob/master/pytorch_segmentation_detection/utils/flops_benchmark.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

This file has been truncated. show original

Example of usage:

github.com

warmspringwinds/pytorch-segmentation-detection/blob/d5df5e066fe9c6078d38b26527d93436bf869b1c/pytorch_segmentation_detection/recipes/pascal_voc/segmentation/flops_counter.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",

This file has been truncated. show original