mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
727ee853b4
Codemodded via `torchfix . --select=TOR203 --fix`. This is a step to unblock https://github.com/pytorch/pytorch/pull/141076 Pull Request resolved: https://github.com/pytorch/pytorch/pull/143691 Approved by: https://github.com/malfet
28 lines
764 B
Python
28 lines
764 B
Python
from torchvision import models
|
|
|
|
import torch
|
|
import torch.autograd.profiler as profiler
|
|
|
|
|
|
for with_cuda in [False, True]:
|
|
model = models.resnet18()
|
|
inputs = torch.randn(5, 3, 224, 224)
|
|
sort_key = "self_cpu_memory_usage"
|
|
if with_cuda and torch.cuda.is_available():
|
|
model = model.cuda()
|
|
inputs = inputs.cuda()
|
|
sort_key = "self_cuda_memory_usage"
|
|
print("Profiling CUDA Resnet model")
|
|
else:
|
|
print("Profiling CPU Resnet model")
|
|
|
|
with profiler.profile(profile_memory=True, record_shapes=True) as prof:
|
|
with profiler.record_function("root"):
|
|
model(inputs)
|
|
|
|
print(
|
|
prof.key_averages(group_by_input_shape=True).table(
|
|
sort_by=sort_key, row_limit=-1
|
|
)
|
|
)
|