torchprof

A minimal dependency library for layer-by-layer profiling of Pytorch models.

All metrics are derived using the PyTorch autograd profiler.

Quickstart

pip install torchprof

import torch import torchvision import torchprof model = torchvision.models.alexnet(pretrained=False).cuda() x = torch.rand([1, 3, 224, 224]).cuda() with torchprof.Profile(model, use_cuda=True) as prof:     model(x) print(prof.display(show_events=False)) # equivalent to `print(prof)` and `print(prof.display())`

Module         | Self CPU total | CPU total | CUDA total ---------------|----------------|-----------|----------- AlexNet        |                |           | ├── features   |                |           | │├── 0         |        1.956ms |   7.714ms |    7.787ms │├── 1         |       68.880us |  68.880us |   69.632us │├── 2         |       85.639us | 155.948us |  155.648us │├── 3         |      253.419us | 970.386us |    1.747ms │├── 4         |       18.919us |  18.919us |   19.584us │├── 5         |       30.910us |  54.900us |   55.296us │├── 6         |      132.839us | 492.367us |  652.192us │├── 7         |       17.990us |  17.990us |   18.432us │├── 8         |       87.219us | 310.776us |  552.544us │├── 9         |       17.620us |  17.620us |   17.536us │├── 10        |       85.690us | 303.120us |  437.248us │├── 11        |       17.910us |  17.910us |   18.400us │└── 12        |       29.239us |  51.488us |   52.288us ├── avgpool    |       49.230us |  85.740us |   88.960us └── classifier |                |           |  ├── 0         |      626.236us |   1.239ms |    1.362ms  ├── 1         |      235.669us | 235.669us |  635.008us  ├── 2         |       17.990us |  17.990us |   18.432us  ├── 3         |       31.890us |  56.770us |   57.344us  ├── 4         |       39.280us |  39.280us |  212.128us  ├── 5         |       16.800us |  16.800us |   17.600us  └── 6         |       38.459us |  38.459us |   79.872us

To see the low level operations that occur within each layer, print the contents of prof.display(show_events=True).

Module                        | Self CPU total | CPU total | CUDA total ------------------------------|----------------|-----------|----------- AlexNet                       |                |           | ├── features                  |                |           | │├── 0                        |                |           | ││├── conv2d                  |       15.740us |   1.956ms |    1.972ms ││├── convolution             |       12.000us |   1.940ms |    1.957ms ││├── _convolution            |       36.590us |   1.928ms |    1.946ms ││├── contiguous              |        6.600us |   6.600us |    6.464us ││└── cudnn_convolution       |        1.885ms |   1.885ms |    1.906ms │├── 1                        |                |           | ││└── relu_                   |       68.880us |  68.880us |   69.632us │├── 2                        |                |           | ││├── max_pool2d              |       15.330us |  85.639us |   84.992us ││└── max_pool2d_with_indices |       70.309us |  70.309us |   70.656us │├── 3                        |                |           | ...

The original Pytorch EventList can be returned by calling raw() on the profile instance.

trace, event_lists_dict = prof.raw() print(trace[2]) # Trace(path=('AlexNet', 'features', '0'), leaf=True, module=Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))) print(event_lists_dict[trace[2].path][0])

---------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  --------------- Name                   Self CPU total %   Self CPU total      CPU total %        CPU total     CPU time avg     CUDA total %       CUDA total    CUDA time avg  Number of Calls ---------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  --------------- conv2d                           0.80%         15.740us          100.00%          1.956ms          1.956ms           25.32%          1.972ms          1.972ms                1 convolution                      0.61%         12.000us           99.20%          1.940ms          1.940ms           25.14%          1.957ms          1.957ms                1 _convolution                     1.87%         36.590us           98.58%          1.928ms          1.928ms           24.99%          1.946ms          1.946ms                1 contiguous                       0.34%          6.600us            0.34%          6.600us          6.600us            0.08%          6.464us          6.464us                1 cudnn_convolution               96.37%          1.885ms           96.37%          1.885ms          1.885ms           24.47%          1.906ms          1.906ms                1 ---------------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  ---------------  --------------- Self CPU time total: 1.956ms CUDA time total: 7.787ms

Layers can be selected for individually using the optional paths kwarg. Profiling is ignored for all other layers.

model = torchvision.models.alexnet(pretrained=False) x = torch.rand([1, 3, 224, 224]) # Layer does not have to be a leaf layer paths = [("AlexNet", "features", "3"), ("AlexNet", "classifier")] with torchprof.Profile(model, paths=paths) as prof:     model(x) print(prof)

Module         | Self CPU total | CPU total | CUDA total ---------------|----------------|-----------|----------- AlexNet        |                |           |            ├── features   |                |           |            │├── 0         |                |           |            │├── 1         |                |           |            │├── 2         |                |           |            │├── 3         |        2.846ms |  11.368ms |    0.000us │├── 4         |                |           |            │├── 5         |                |           |            │├── 6         |                |           |            │├── 7         |                |           |            │├── 8         |                |           |            │├── 9         |                |           |            │├── 10        |                |           |            │├── 11        |                |           |            │└── 12        |                |           |            ├── avgpool    |                |           |            └── classifier |       12.016ms |  12.206ms |    0.000us  ├── 0         |                |           |             ├── 1         |                |           |             ├── 2         |                |           |             ├── 3         |                |           |             ├── 4         |                |           |             ├── 5         |                |           |             └── 6         |                |           |

• Self CPU Time vs CPU Time

LICENSE

MIT