Function hook

PATH/helper/param_count.py:83–88

(module, x)

    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            def hook(module, x):
                out_h = int((x[0].size()[2] + 2 * module.padding[0] - module.kernel_size[0]) / module.stride[0] + 1)
                out_w = int((x[0].size()[3] + 2 * module.padding[1] - module.kernel_size[1]) / module.stride[1] + 1)
                ops = module.in_channels * module.out_channels * module.kernel_size[0] * module.kernel_size[1] * out_h * out_w / module.groups * multi_add
                flop_counts.append(ops)
                param_counts.append(get_layer_param(module))
            m.register_forward_pre_hook(hook)

        elif isinstance(m, nn.ReLU) or isinstance(m, nn.PReLU):