How to measure GFLOPs

[Youngwan Lee]

Hi,

I found how to measure model size, but couldn't find the calculation of the GFLOPs.

Could you let me know how to measure the GFLOPs?

Thanks in advance :)

[Fan Bao]

We use the facebookresearch/fvcore lib. Here is a demo to calculate the GFLOPs of U-ViT-H on ImageNet 256x256

from fvcore.nn import FlopCountAnalysis
from libs.uvit import UViT
model = UViT(
  img_size=32,
  patch_size=2,
  in_chans=4,
  embed_dim=1152,
  depth=28,
  num_heads=16,
  mlp_ratio=4,
  qkv_bias=False,
  mlp_time_embed=False,
  num_classes=1001,
  use_checkpoint=True,
  conv=False
)

inputs = (torch.rand(1, 4, 32, 32), torch.tensor([1]), torch.tensor([1]))
flops = FlopCountAnalysis(model, inputs)
print("FLOPS of imagenet256_uvit_huge: ", f'{flops.total() / 1e9}G')

[Youngwan Lee]

@baofff Thanks :)