Measure neural network device specific metrics (latency, flops, etc.)
Project description
torchprof
Attention! This library is deprecated due to the PyTorch 1.9 changes to the torch profiler. Please use the official profiler. Thank you!
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()
# `profile_memory` was added in PyTorch 1.6, this will output a runtime warning if unsupported.
with torchprof.Profile(model, use_cuda=True, profile_memory=True) as prof:
model(x)
# equivalent to `print(prof)` and `print(prof.display())`
print(prof.display(show_events=False))
Module | Self CPU total | CPU total | Self CUDA total | CUDA total | Self CPU Mem | CPU Mem | Self CUDA Mem | CUDA Mem | Number of Calls
---------------|----------------|-----------|-----------------|------------|--------------|---------|---------------|-----------|----------------
AlexNet | | | | | | | | |
├── features | | | | | | | | |
│├── 0 | 1.832ms | 7.264ms | 1.831ms | 7.235ms | 0 b | 0 b | 756.50 Kb | 3.71 Mb | 1
│├── 1 | 51.858us | 76.564us | 51.296us | 76.896us | 0 b | 0 b | 0 b | 0 b | 1
│├── 2 | 75.993us | 157.855us | 77.600us | 145.184us | 0 b | 0 b | 547.00 Kb | 1.60 Mb | 1
│├── 3 | 263.526us | 1.142ms | 489.472us | 1.918ms | 0 b | 0 b | 547.00 Kb | 2.68 Mb | 1
│├── 4 | 28.824us | 41.197us | 28.672us | 43.008us | 0 b | 0 b | 0 b | 0 b | 1
│├── 5 | 55.264us | 120.016us | 55.200us | 106.400us | 0 b | 0 b | 380.50 Kb | 1.11 Mb | 1
│├── 6 | 175.591us | 681.011us | 212.896us | 818.080us | 0 b | 0 b | 253.50 Kb | 8.27 Mb | 1
│├── 7 | 27.622us | 39.494us | 26.848us | 39.296us | 0 b | 0 b | 0 b | 0 b | 1
│├── 8 | 140.204us | 537.162us | 204.832us | 781.280us | 0 b | 0 b | 169.00 Kb | 10.20 Mb | 1
│├── 9 | 27.532us | 39.364us | 26.816us | 39.136us | 0 b | 0 b | 0 b | 0 b | 1
│├── 10 | 138.621us | 530.929us | 171.008us | 650.432us | 0 b | 0 b | 169.00 Kb | 7.08 Mb | 1
│├── 11 | 27.712us | 39.645us | 27.648us | 39.936us | 0 b | 0 b | 0 b | 0 b | 1
│└── 12 | 54.813us | 118.823us | 55.296us | 107.360us | 0 b | 0 b | 108.00 Kb | 324.00 Kb | 1
├── avgpool | 58.329us | 116.577us | 58.368us | 111.584us | 0 b | 0 b | 36.00 Kb | 108.00 Kb | 1
└── classifier | | | | | | | | |
├── 0 | 79.169us | 167.495us | 78.848us | 145.408us | 0 b | 0 b | 45.00 Kb | 171.00 Kb | 1
├── 1 | 404.070us | 423.755us | 793.600us | 793.600us | 0 b | 0 b | 16.00 Kb | 32.00 Kb | 1
├── 2 | 30.097us | 43.512us | 29.792us | 43.904us | 0 b | 0 b | 0 b | 0 b | 1
├── 3 | 53.390us | 121.042us | 53.248us | 99.328us | 0 b | 0 b | 20.00 Kb | 76.00 Kb | 1
├── 4 | 64.622us | 79.902us | 236.544us | 236.544us | 0 b | 0 b | 16.00 Kb | 32.00 Kb | 1
├── 5 | 28.854us | 41.067us | 28.544us | 41.856us | 0 b | 0 b | 0 b | 0 b | 1
└── 6 | 62.258us | 77.356us | 95.232us | 95.232us | 0 b | 0 b | 4.00 Kb | 8.00 Kb | 1
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 | Self CUDA total | CUDA total | Self CPU Mem | CPU Mem | Self CUDA Mem | CUDA Mem | Number of Calls
------------------------------------|----------------|-----------|-----------------|------------|--------------|---------|---------------|-----------|----------------
AlexNet | | | | | | | | |
├── features | | | | | | | | |
│├── 0 | | | | | | | | |
││├── aten::conv2d | 15.630us | 1.832ms | 14.176us | 1.831ms | 0 b | 0 b | 0 b | 756.50 Kb | 1
││├── aten::convolution | 9.768us | 1.816ms | 9.056us | 1.817ms | 0 b | 0 b | 0 b | 756.50 Kb | 1
││├── aten::_convolution | 45.005us | 1.807ms | 34.432us | 1.808ms | 0 b | 0 b | 0 b | 756.50 Kb | 1
││├── aten::contiguous | 8.738us | 8.738us | 8.480us | 8.480us | 0 b | 0 b | 0 b | 0 b | 3
││├── aten::cudnn_convolution | 1.647ms | 1.683ms | 1.745ms | 1.750ms | 0 b | 0 b | -18.00 Kb | 756.50 Kb | 1
││├── aten::empty | 21.249us | 21.249us | 0.000us | 0.000us | 0 b | 0 b | 774.50 Kb | 774.50 Kb | 2
││├── aten::resize_ | 7.635us | 7.635us | 0.000us | 0.000us | 0 b | 0 b | 0 b | 0 b | 2
││├── aten::stride | 1.902us | 1.902us | 0.000us | 0.000us | 0 b | 0 b | 0 b | 0 b | 4
││├── aten::reshape | 6.081us | 17.833us | 2.048us | 2.048us | 0 b | 0 b | 0 b | 0 b | 1
││├── aten::view | 11.752us | 11.752us | 0.000us | 0.000us | 0 b | 0 b | 0 b | 0 b | 1
││└── aten::add_ | 57.248us | 57.248us | 18.432us | 18.432us | 0 b | 0 b | 0 b | 0 b | 1
│├── 1 | | | | | | | | |
││├── aten::relu_ | 27.152us | 51.858us | 25.696us | 51.296us | 0 b | 0 b | 0 b | 0 b | 1
││└── aten::threshold_ | 24.706us | 24.706us | 25.600us | 25.600us | 0 b | 0 b | 0 b | 0 b | 1
│├── 2 | | | | | | | | |
...
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 % Self CPU CPU total % CPU total CPU time avg Self CUDA Self CUDA % CUDA total CUDA time avg CPU Mem Self CPU Mem CUDA Mem Self CUDA Mem # of Calls
--------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
aten::conv2d 0.85% 15.630us 100.00% 1.832ms 1.832ms 14.176us 0.77% 1.831ms 1.831ms 0 b 0 b 756.50 Kb 0 b 1
aten::convolution 0.53% 9.768us 99.15% 1.816ms 1.816ms 9.056us 0.49% 1.817ms 1.817ms 0 b 0 b 756.50 Kb 0 b 1
aten::_convolution 2.46% 45.005us 98.61% 1.807ms 1.807ms 34.432us 1.88% 1.808ms 1.808ms 0 b 0 b 756.50 Kb 0 b 1
aten::contiguous 0.20% 3.707us 0.20% 3.707us 3.707us 3.680us 0.20% 3.680us 3.680us 0 b 0 b 0 b 0 b 1
aten::cudnn_convolution 89.90% 1.647ms 91.86% 1.683ms 1.683ms 1.745ms 95.27% 1.750ms 1.750ms 0 b 0 b 756.50 Kb -18.00 Kb 1
aten::empty 0.66% 12.102us 0.66% 12.102us 12.102us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 756.50 Kb 756.50 Kb 1
aten::contiguous 0.15% 2.706us 0.15% 2.706us 2.706us 2.560us 0.14% 2.560us 2.560us 0 b 0 b 0 b 0 b 1
aten::resize_ 0.39% 7.164us 0.39% 7.164us 7.164us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::contiguous 0.13% 2.325us 0.13% 2.325us 2.325us 2.240us 0.12% 2.240us 2.240us 0 b 0 b 0 b 0 b 1
aten::resize_ 0.03% 0.471us 0.03% 0.471us 0.471us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::stride 0.06% 1.092us 0.06% 1.092us 1.092us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::stride 0.02% 0.280us 0.02% 0.280us 0.280us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::stride 0.01% 0.270us 0.01% 0.270us 0.270us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::stride 0.01% 0.260us 0.01% 0.260us 0.260us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::empty 0.50% 9.147us 0.50% 9.147us 9.147us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 18.00 Kb 18.00 Kb 1
aten::reshape 0.33% 6.081us 0.97% 17.833us 17.833us 2.048us 0.11% 2.048us 2.048us 0 b 0 b 0 b 0 b 1
aten::view 0.64% 11.752us 0.64% 11.752us 11.752us 0.000us 0.00% 0.000us 0.000us 0 b 0 b 0 b 0 b 1
aten::add_ 3.12% 57.248us 3.12% 57.248us 57.248us 18.432us 1.01% 18.432us 18.432us 0 b 0 b 0 b 0 b 1
--------------------------- ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------ ------------
Self CPU time total: 1.832ms
CUDA time total: 1.831ms
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 | Number of Calls
---------------|----------------|-----------|----------------
AlexNet | | |
├── features | | |
│├── 0 | | |
│├── 1 | | |
│├── 2 | | |
│├── 3 | 3.162ms | 12.626ms | 1
│├── 4 | | |
│├── 5 | | |
│├── 6 | | |
│├── 7 | | |
│├── 8 | | |
│├── 9 | | |
│├── 10 | | |
│├── 11 | | |
│└── 12 | | |
├── avgpool | | |
└── classifier | 11.398ms | 12.130ms | 1
├── 0 | | |
├── 1 | | |
├── 2 | | |
├── 3 | | |
├── 4 | | |
├── 5 | | |
└── 6 | | |
Citation
If this software is useful to your research, I would greatly appreciate a citation in your work.
@misc{awwong1-torchprof,
title = {torchprof},
author = {Alexander William Wong},
month = 12,
year = 2020,
url = {https://github.com/awwong1/torchprof}
note = {https://github.com/awwong1/torchprof}
}
LICENSE
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
torchprof-1.4.0.tar.gz
(13.4 kB
view hashes)
Built Distribution
torchprof-1.4.0-py3-none-any.whl
(11.4 kB
view hashes)
Close
Hashes for torchprof-1.4.0-py3-none-any.whl
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 | 2fbc2d4a7004f3db02599e5d9788029fa02094832ebf14d02a9db50bc7aa5d15 |
|
| MD5 | ff10eb7688979aae84f91f20f4b7c85c |
|
| BLAKE2b-256 | 97dbf6469e51e99d1885336833227987d82160384b98f7cf56ab671ebbd9e346 |
