NAFNet

NAFNet copied to clipboard

hi, Thanks first for the excellent work! And can you share the script of counting model MACs?

Hi, knsong, Thanks!

We use an open-source tool called ptflops An example of one NAFBlock results like this:

(module): NAFBlock(
          0.466 M, 0.686% Params, 1.639 GMac, 2.588% MACs,
          (conv1): Conv2d(0.132 M, 0.194% Params, 0.539 GMac, 0.851% MACs, 256, 512, kernel_size=(1, 1), stride=(1, 1))
          (conv2): Conv2d(0.005 M, 0.008% Params, 0.021 GMac, 0.033% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=512)
          (conv3): Conv2d(0.066 M, 0.097% Params, 0.269 GMac, 0.425% MACs, 256, 256, kernel_size=(1, 1), stride=(1, 1))
          (sca): Sequential(
            0.066 M, 0.097% Params, 0.001 GMac, 0.002% MACs,
            (0): AdaptiveAvgPool2d(0.0 M, 0.000% Params, 0.001 GMac, 0.002% MACs, output_size=1)
            (1): Conv2d(0.066 M, 0.097% Params, 0.0 GMac, 0.000% MACs, 256, 256, kernel_size=(1, 1), stride=(1, 1))
          )
          (sg): SimpleGate(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, )
          (conv4): Conv2d(0.132 M, 0.194% Params, 0.539 GMac, 0.851% MACs, 256, 512, kernel_size=(1, 1), stride=(1, 1))
          (conv5): Conv2d(0.066 M, 0.097% Params, 0.269 GMac, 0.425% MACs, 256, 256, kernel_size=(1, 1), stride=(1, 1))
          (norm1): LayerNorm2d(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, )
          (norm2): LayerNorm2d(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, )
        )

As you can notice, it ignores the custom layers we implemented, i.e. SimpleGate && LayerNorm2d, we simply add the flops of these components manually. In fact, these components only occupy ~0.5% of the total MACs. So if you only want an approximate number, you can directly ignore the calculation amount of these components.