Pytorch-UNet icon indicating copy to clipboard operation
Pytorch-UNet copied to clipboard

How to accurately calculate Model Parameters and GFLOPs?

Open RobertLee0522 opened this issue 1 year ago • 0 comments

I'm using the following method to calculate parameters and GFLOPs, with an image size of 840x420, img_scale=0.5, and batch size set to 1:

  • Image size: 840x420
  • img_scale: 0.5
  • Batch size: 1
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6  # Convert to million (M) units

def calculate_gflops(model, input_tensor):
    flops, params = profile(model, inputs=(input_tensor, ))
    return flops ,params

input_tensor = torch.randn((1, 3, 420, 240))  
num_params = count_parameters(model)
gflops = calculate_gflops(model, input_tensor)

print(f"Number of parameters: {num_params} M")
print(f"GFLOPs: {gflops[0]/1000**3}G")
print(f"Params: {gflops[1]/1000**2}M")

Results:

  • Number of parameters: 31.037698 M
  • GFLOPs: 83.922432G
  • Params: 31.037698M

Is this approach correct for estimating the model's parameters and GFLOPs? I would appreciate any feedback or suggestions. Thank you!

RobertLee0522 avatar Jan 26 '24 03:01 RobertLee0522