How to use the thop.profile.profile function in thop
To help you get started, we’ve selected a few thop examples, based on popular ways it is used in public projects.
Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
Lyken17 / pytorch-OpCounter / benchmark / evaluate_famous_models.py View on Github 
and callable(models.__dict__[name]))
print("%s | %s | %s" % ("Model", "Params(M)", "FLOPs(G)"))
print("---|---|---")
device = "cpu"
if torch.cuda.is_available():
device = "cuda"
for name in model_names:
model = models.__dict__[name]().to(device)
dsize = (1, 3, 224, 224)
if "inception" in name:
dsize = (1, 3, 299, 299)
inputs = torch.randn(dsize).to(device)
total_ops, total_params = profile(model, (inputs,), verbose=False)
print("%s | %.2f | %.2f" % (name, total_params / (1000 ** 2), total_ops / (1000 ** 3)))thop
A tool to count the FLOPs of PyTorch model.
