pytorch

1
import torchvision.models as models
2

3
import torch
4
import torch.autograd.profiler as profiler
5

6

7
for with_cuda in [False, True]:
8
    model = models.resnet18()
9
    inputs = torch.randn(5, 3, 224, 224)
10
    sort_key = "self_cpu_memory_usage"
11
    if with_cuda and torch.cuda.is_available():
12
        model = model.cuda()
13
        inputs = inputs.cuda()
14
        sort_key = "self_cuda_memory_usage"
15
        print("Profiling CUDA Resnet model")
16
    else:
17
        print("Profiling CPU Resnet model")
18

19
    with profiler.profile(profile_memory=True, record_shapes=True) as prof:
20
        with profiler.record_function("root"):
21
            model(inputs)
22

23
    print(
24
        prof.key_averages(group_by_input_shape=True).table(
25
            sort_by=sort_key, row_limit=-1
26
        )
27
    )
28