October 19, 2024

It is very important to know structures of neural networks, and how many parameters the networks have in order to know the capabilities of the models and how much computational resources are required.

from torch.profiler import profile, record_function, ProfilerActivity
with profile(activities=[ProfilerActivity.CPU], record_shapes=True) as prof:
    with record_function('model_inference'):
        result = model(image)