Profile a model layer by layer

1

Hey, I find it difficult to understand wether it is possible to profile a pytorch layerwise.

So far, I wrap the layers in my model (nn.Module) using

with profiler.record_function("layer_name_shown_in_summary"):
     self.layer= SomeLayer(...)

I wrap the optimization loop in this expression

with torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                    torch.profiler.ProfilerActivity.CUDA,  # Optional: if using GPU
                ],
                record_shapes=True,
                with_stack=True,
                with_modules=True
            ) as prof:
       do the training here

and I print out the statistics using this comand here

print(prof.key_averages(group_by_input_shape=True).table(
        sort_by   = "self_cuda_time_total",
        row_limit = 20
        )
    )

The result looks like this

Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls    Input Shapes
...
aten::mul         0.01%     245.862us         0.04%     899.349us      49.964us      57.740ms         3.98%      57.740ms       3.208ms            18     [[SomeBatchDim, SomeVectorDim], [SomeVectorDim]]

My question is:

  • Is the time recorded layerwise?
  • Can I group and summarize the single function for each layer?

Any help would be much appreciated.

2

The output should contain the record_function tag as also seen in the profiler tutorial.