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.
I have some questions:

Is it normal to include flops of ReLU, Batch normalization, …?
It seems common to consider the spatial dimension. For example, when calculating Conv2d layer, I need to know the image size. What is the common size of the original image size of width and height?
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 ReLU, we don’t know the previous state.  )

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 -> 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') >= 0:
        C1 = int(params[0])
        C2 = int(params[1])
        flops = C1*C2
        
    elif layer.find('Conv2d') >= 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:
        >>> net = torchvision.models.resnet18()
        >>> 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