RoyChao19477 / SEMamba

This is the official implementation of the SEMamba paper. (Accepted to IEEE SLT 2024)
118 stars 11 forks source link

FLOPS calculation #8

Open dazz-debug56 opened 1 month ago

dazz-debug56 commented 1 month ago

Thank you for the great work.

I would like to ask how did you compute the FLOPs of the model and whether you can share the method?
fvcore doesn't seem to work here due to the Mamba blocks.

RoyChao19477 commented 3 weeks ago

Hi dazz-debug56,

Thank you for your interest. We used the thop library to calculate the FLOPs. Below is the code snippet we used:

from thop import profile

flops, param = profile(generator, inputs=(noisy_mag, noisy_pha))
print('FLOPs = {:.3f} GFLOPs'.format(flops / 1e9))

This calculation was performed with an input audio length of 2 seconds and a batch size of 1 on a single GPU.

Best regards, Roy Chao