Memory Profile Results

username: Woosung
cooked: <p>Hi everyone,</p>
<p>I have a question regarding the memory profiling results of ResNet.</p>
<p>I followed the tutorial from the <a href="https://pytorch.org/blog/understanding-gpu-memory-1/" rel="noopener nofollow ugc">PyTorch blog</a>. The main difference in my case is that I profiled the memory usage during the inference step, rather than training.</p>
<p>I profiled the model-building process and 4 iterations of inference.<br>
Below is a snapshot of the memory usage visualization.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e.png" data-download-href="https://discuss.pytorch.org/uploads/default/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e_2_517x318.png" alt="image" data-base62-sha1="jMF0lHTGtxRqOOOR9xpZJRD0BEW" width="517" height="318" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e_2_517x318.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e_2_775x477.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8aa97b8983fb747b7c86583fffd3f1600f0ebd0e_2_1034x636.png 2x" data-dominant-color="DCCCC0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1736×1068 84 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My question is: why does the memory footprint of the first iteration (inference) appear significantly smaller compared to subsequent iterations?</p>
<p>Below is my full code for reference.</p>
<pre data-code-wrap="py"><code class="lang-py"># (c) Meta Platforms, Inc. and affiliates. 
import logging
import socket
from datetime import datetime, timedelta

import torch

from torch.autograd.profiler import record_function
from torchvision import models

logging.basicConfig(
    format="%(levelname)s:%(asctime)s %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S",
)

logger: logging.Logger = logging.getLogger(__name__)
logger.setLevel(level=logging.INFO)

TIME_FORMAT_STR: str = "%b_%d_%H_%M_%S"

def trace_handler(prof: torch.profiler.profile):
    # Prefix for file names.
    host_name = socket.gethostname()
    timestamp = datetime.now().strftime(TIME_FORMAT_STR)
    file_prefix = f"{host_name}_{timestamp}"

    # Construct the trace file.
    prof.export_chrome_trace(f"{file_prefix}.json.gz")

    # Construct the memory timeline file.
    prof.export_memory_timeline(f"{file_prefix}.html", device="cuda:0")

def run_resnet50(num_iters=4, device="cuda:0"):
    with torch.profiler.profile(
        activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
        ],
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
        on_trace_ready=trace_handler,
        with_modules=True,
    ) as prof:
        with record_function("## Prepare ##"):
            model = models.resnet50().to(device=device)
            inputs = torch.randn(1, 3, 224, 224, device=device)
            model.eval()
        for _ in range(num_iters):
            with record_function("## forward ##"):
                pred = model(inputs)

if __name__ == "__main__":
    # Warm up
    run_resnet50()
    # Run the resnet50 model
    run_resnet50()
</code></pre>
<p>Thanks in advance!</p>

username: ptrblck
cooked: <p>You are not wrapping your code into a <code>torch.no_grad()</code> context so intermediate activations are kept alive (without a need for these). The second iterations is thus allocating these activations for the second forward pass while the activations from the first forward pass are still stored and increases the memory usage further (this would be my guess without digging into it).</p>