csguoh / MambaIR

[ECCV2024] An official pytorch implement of the paper "MambaIR: A simple baseline for image restoration with state-space model".
Apache License 2.0
453 stars 41 forks source link

param,flops #44

Open wzylove opened 4 months ago

wzylove commented 4 months ago

Hello, author, first of all, thank you very much for your work. I recently wanted to make some improvements in lightweight. I used the following code, but I found that an error occurred when testing your code. Can you provide the code for testing flops and param? image

csguoh commented 4 months ago

Hi, thx for your interest in this work.

The code for calculating MACs and #param is actually in this floder: /analysis/flops_param.py. However, it should be noted that the params obtained from some third-party liarbries, such as thop, may ignore certain parts of the model, due to the BUG of these liarbries. The bug can lead to inaccurate counts which is also revealed in original Mamba repo. To obtain the correct #param, you may use this line to obtain the params:

trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)

The obtained results are larger than the reported one in the paper. We have recently planned to train another MambaIR-light version under the correct parameter counts setup. Stay tuned.

zhousai-zs commented 4 months ago

Hi, thx for your interest in this work.

The code for calculating MACs and #param is actually in this floder: /analysis/flops_param.py. However, it should be noted that the params obtained from some third-party liarbries, such as thop, may ignore certain parts of the model, due to the BUG of these liarbries. The bug can lead to inaccurate counts which is also revealed in original Mamba repo. To obtain the correct #param, you may use this line to obtain the params:

trainable_num = sum(p.numel() for p in model.parameters() if p.requires_grad)

The obtained results are larger than the reported one in the paper. We have recently planned to train another MambaIR-light version under the correct parameter counts setup. Stay tuned.

So, the FLOPs obtained from ./analysis/flops_param.py is right? I have try to calc the FLOPs of SS2D,and I reduce K sequence to 1 sequence, but gain the same FLOPs and Parameters. In my opinion, it is a wrong answer.

csguoh commented 3 months ago

Hi, the FLOPs also encounters the same problem when using some third-party libraries. To obtain the accurate FLOPs, you may refer this issue in the Vmamba repo, and register the missed params in the fvcore. We are working on the accurate calculation and will release the revised code soon.

zhousai-zs commented 3 months ago

Hi, the FLOPs also encounters the same problem when using some third-party libraries. To obtain the accurate FLOPs, you may refer this issue in the Vmamba repo, and register the missed params in the fvcore. We are working on the accurate calculation and will release the revised code soon.

OK, Thanks for your reply, and looking forward to your update about calculating FLOPs, Params and inference time.