def on_start(self, task) -> None:
"""Profile the forward pass."""
logging.info("Profiling forward pass...")
batchsize_per_replica = task.get_batchsize_per_replica()
input_shape = task.base_model.input_shape
p = profile(
task.model,
batchsize_per_replica=batchsize_per_replica,
input_shape=input_shape,
input_key=getattr(task.base_model, "input_key", None),
)
logging.info(summarize_profiler_info(p))
def on_start(self, task: "tasks.ClassyTask") -> None:
"""Profile the forward pass."""
logging.info("Profiling forward pass...")
batchsize_per_replica = getattr(
task.dataloaders[task.phase_type].dataset, "batchsize_per_replica",
1)
input_shape = task.base_model.input_shape
p = profile(
task.model,
batchsize_per_replica=batchsize_per_replica,
input_shape=input_shape,
input_key=getattr(task.base_model, "input_key", None),
)
logging.info(summarize_profiler_info(p))