milesial/Pytorch-UNet

How to accurately calculate Model Parameters and GFLOPs?

RobertLee0522 opened this issue · 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!