# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.9.7 64-bit (''mlops'': conda)'
#     language: python
#     name: python3
# ---

import torch
import torchvision.models as models
from torch.profiler import profile, record_function, ProfilerActivity, tensorboard_trace_handler

model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)

with profile(activities=[ProfilerActivity.CPU], record_shapes=True, profile_memory=True,on_trace_ready=tensorboard_trace_handler("profiles/")) as prof:
   with record_function("model_inference"):
      model(inputs)

print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10))

print(prof.key_averages(group_by_input_shape=True).table(sort_by="cpu_time_total", row_limit=30))

prof.export_chrome_trace("trace.json")



StarcoderdataJupyterScriptsDedupFiltered