본문 바로가기

Deep Learning

How to count the number of parameters of a NN and measure FLOPs required for the NN

반응형
from fvcore.nn import FlopCountAnalysis
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

 

ResNet = ConvNet(use_pretrained=True,
                 feature_extract=False,
                 resent_model=saved_args.resnet_model)
N_param = count_parameters(ResNet) / 1e6
input_tensor = torch.randn(1, 3, 320, 640)
flops = FlopCountAnalysis(ResNet, input_tensor)
print(f'>> ConvNet Model is loaded .. : {N_param} M, {flops.total() / 1e9} FLOPs')