pytorch框架下计算视频工作的flops与params（计算网络计算量和参数量））

这个工作做到的话可以显示出flops（FLOPS：全称是floating point operations per second），parameters（参数，网络中对应的参数）。

#coding:utf8
import torch
import torchvision

import torch.nn as nn
from torch.autograd import Variable
import torchvision.models as models

import numpy as np
##usage: add to train.py or test.py: misc.print_model_parm_nums(model)
##  misc.print_model_parm_flops(model,inputs)
def print_model_parm_nums(model):
    total = sum([param.nelement() for param in model.parameters()])
    print('  + Number of params: %.2f(e6)' % (total / 1e6))

def print_model_parm_flops(model,input):

    # prods = {}
    # def save_prods(self, input, output):
        # print 'flops:{}'.format(self.__class__.__name__)
        # print 'input:{}'.format(input)
        # print '_dim:{}'.format(input[0].dim())
        # print 'input_shape:{}'.format(np.prod(input[0].shape))
        # grads.append(np.prod(input[0].shape))

    prods = {}
    def save_hook(name):
        def hook_per(self, input, output):
            # print 'flops:{}'.format(self.__class__.__name__)
            # print 'input:{}'.format(input)
            # print '_dim:{}'.format(input[0].dim())
            # print 'input_shape:{}'.format(np.prod(input[0].shape))
            # prods.append(np.prod(input[0].shape))
            prods[name] = np.prod(input[0].shape)
            # prods.append(np.prod(input[0].shape))
        return hook_per

    list_1=[]
    def simple_hook(self, input, output):
        list_1.append(np.prod(input[0].shape))
    list_2={}
    def simple_hook2(self, input, output):
        list_2['names'] = np.prod(input[0].shape)


    multiply_adds = False
    list_conv=[]
    def conv_hook(self, input, output):
        batch_size, input_channels, input_time, input_height, input_width = input[0].size()
        output_channels, output_time, output_height, output_width = output[0].size()
        
        kernel_ops = self.kernel_size[0] * self.kernel_size[1] * self.kernel_size[2] *(self.in_channels / self.groups) * (2 if multiply_adds else 1)
        bias_ops = 1 if self.bias is not None else 0

        params = output_channels * (kernel_ops + bias_ops)
        flops = batch_size * params  * output_time * output_height * output_width
        list_conv.append(flops)


    list_linear=[] 
    def linear_hook(self, input, output):
        batch_size = input[0].size(0) if input[0].dim() == 2 else 1

        weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
        bias_ops = self.bias.nelement()

        flops = batch_size * (weight_ops + bias_ops)
        list_linear.append(flops)

    list_bn=[] 
    def bn_hook(self, input, output):
        list_bn.append(input[0].nelement())
    
    list_fc=[] 
    def fc_hook(self, input, output):
        list_bn.append(input[0].nelement())


    list_relu=[] 
    def relu_hook(self, input, output):
        list_relu.append(input[0].nelement())

    list_pooling=[]
    #def pooling_hook(self, input, output):
     #   batch_size, input_channels, input_time,input_height, input_width = input[0].size()
      #  output_channels, output_time, output_height, output_width = output[0].size()

       # kernel_ops = self.kernel_size * self.kernel_size*self.kernel_size
        #bias_ops = 0
        #params = output_channels * (kernel_ops + bias_ops)
        #flops = batch_size * params * output_height * output_width * output_time

        #list_pooling.append(flops)


            
    def foo(net):
        childrens = list(net.children())
        if not childrens:
            if isinstance(net, torch.nn.Conv3d):
                # net.register_forward_hook(save_hook(net.__class__.__name__))
                # net.register_forward_hook(simple_hook)
                # net.register_forward_hook(simple_hook2)
                net.register_forward_hook(conv_hook)
            if isinstance(net, torch.nn.Linear):
                net.register_forward_hook(linear_hook)
            if isinstance(net, torch.nn.BatchNorm3d):
                net.register_forward_hook(bn_hook)
            if isinstance(net, torch.nn.ReLU):
                net.register_forward_hook(relu_hook)
            if isinstance(net, torch.nn.ReLU):
                net.register_forward_hook(relu_hook)
            #if isinstance(net, torch.nn.MaxPool3d) or isinstance(net, torch.nn.AvgPool2d):
             #   net.register_forward_hook(pooling_hook)
            return
        for c in childrens:
                foo(c)

    foo(model)
    output = model(input)
    total_flops = (sum(list_conv)+sum(list_linear))#+sum(list_bn)+sum(list_relu))
    print('  + Number of FLOPs: %.5f(e9)' % (total_flops / 1e9))

上面代码经测试已经适配到1.0版本，用法直接定义上部分两个函数即可.
问题已经解决，因为当时我们跑的是8帧的，所以说相对的time这一块会有所变化，也就是说实际上的得到的比较大的那个数是因为我们现在用的是32帧，除以4就可以得到之前的对应数字。
实际上算的时候batchsize设置默认为1的时候，并不算整体的数据的计算量。
疑点：

1. 如果是计算maxpooling的时候是否需要加进去？因为maxpooling实际上只是找到区域最大的一块数。
2. 如果是计算interpolation的话应该怎么计算其中的flops？
   基础知识参考这篇文章
   上面这篇讲的已经很好了，对于视频的扩展到conv3d我们也已经把他进行改进了。
   基础知识这里面也有讲到。
   NVIDIA给出的解释也有一些参考价值。