fvcore_pytorch_model_flops_and_parameter_count_analysis.py

1import torch
2import torchvision.models as models
3from fvcore.nn import FlopCountAnalysis, parameter_count_table
4
5# 1. Initialize a model (e.g., a ResNet-18 from torchvision)
6model = models.resnet18()
7
8# 2. Create a dummy input that matches the model's expected input shape
9# ResNet-18 expects (batch_size, channels, height, width)
10inputs = (torch.randn(1, 3, 224, 224),)
11
12# 3. Initialize the FlopCountAnalysis
13flops = FlopCountAnalysis(model, inputs)
14
15# 4. Get the total FLOPs
16print(f"Total FLOPs: {flops.total()}")
17
18# 5. Get FLOPs per operator (returns a counter-like dictionary)
19print("FLOPs per operator:")
20print(flops.by_operator())
21
22# 6. Generate a formatted table of parameter counts
23print("\nParameter Count Table:")
24print(parameter_count_table(model))