import torch from thop import profile def FLOPs_and_Params(model, img_size, device): x = torch.randn(1, 3, img_size, img_size).to(device) print('==============================') flops, params = profile(model, inputs=(x, ), verbose=False) print('==============================') print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2)) print('Params : {:.2f} M'.format(params / 1e6)) if __name__ == "__main__": pass