Counting FLOPS during DropOut about pytorch-opcounter HOT 4 CLOSED

# assume we want to ignore FLOPs in batchnorm layer
def profile(model, input_size, custom_ops={}):
    def add_hooks(m):
        if len(list(m.children())) > 0:
            return

        # ========  add one judgement here ===========
        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
            return 

        m.register_buffer('total_ops', torch.zeros(1))
        m.register_buffer('total_params', torch.zeros(1))

        for p in m.parameters():
            m.total_params += torch.Tensor([p.numel()])

        m_type = type(m)
        fn = None

        if m_type in custom_ops:
            fn = custom_ops[m_type]
        elif m_type in register_hooks:
            fn = register_hooks[m_type]
        else:
            logging.warning("Not implemented for ", m)

        if fn is not None:
            logging.info("Register FLOP counter for module %s" % str(m))
            m.register_forward_hook(fn)

    model.eval()
    model.apply(add_hooks)

    x = torch.zeros(input_size)
    model(x)

    total_ops = 0
    total_params = 0
    for m in model.modules():
        if len(list(m.children())) > 0: # skip for non-leaf module
            continue
        total_ops += m.total_ops
        total_params += m.total_params

    total_ops = total_ops.item()
    total_params = total_params.item()

    return total_ops, total_params

    for m in model.modules():
        if len(list(m.children())) > 0: # skip for non-leaf module
            continue
        # print layer-wise information here.
        print(str(m),  m.total_ops, m.total_params)
        total_ops += m.total_ops
        total_params += m.total_params