Difference in memory usage between 1.7.1 and 1.7.1+cpu

czotti: <p>Hi,</p>
<p>I tried the latest two versions of PyTorch respectively 1.7.1 and 1.7.1+cpu.</p>
<p>I see a really big difference in memory usage for the same model between the two. I run the model on CPU only never allocate anything on GPU (checked with nvidia-smi).<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/d/cd9eb83883f7b085f01213159aff060a53a3f394.png" data-download-href="https://discuss.pytorch.org/uploads/default/cd9eb83883f7b085f01213159aff060a53a3f394" title="Resnet 101"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/d/cd9eb83883f7b085f01213159aff060a53a3f394.png" alt="Resnet 101" data-base62-sha1="tkZW4siRlvDdpnmVCpchdjURmFm" width="690" height="366" data-dominant-color="FBFAFB"><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">Resnet 101</span><span class="informations">698×371 15.1 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><br>
Vertical axis is memory usage, and horizontal axis is input size <code>(1, 3, size, size)</code> where <code>size</code> take three values 224, 256, 512.<br>
I don’t find anything that make this huge difference, since the code is the same. I installed the two packages versions on different virtualenvs.</p>
<p>here is the code to reproduce</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import argparse
from memory_profiler import memory_usage
from torchvision import models
from torch.utils import mkldnn as mkldnn_utils
import time

parser = argparse.ArgumentParser()
parser.add_argument("size", type=int, default=224)
parser.add_argument("iter", type=int, default=100)
parser.add_argument("-m", "--mkldnn", action="store_true")
parser.add_argument("-c", "--cache", action="store_true")

args = parser.parse_args()


def forward(
    net,
    use_mkldnn=False,
    iteration=1,
    weight_cache=False,
    size=224,
):
    net.eval()
    batch = torch.rand(1, 3, size, size)
    if use_mkldnn:
        net = mkldnn_utils.to_mkldnn(net)
        batch = batch.to_mkldnn()
        if weight_cache:
            # using weight cache which will reduce weight reorder
            fname = "test.script.pt"
            traced = torch.jit.trace(net, batch, check_trace=False)
            script = traced.save(fname)
            net = torch.jit.load(fname)

    start_time = time.time()
    for i in range(iteration):
        with torch.no_grad():
            net(batch)
    print(
        f"time: {(time.time() - start_time) / iteration}, iteration: {iteration}"
    )


net = models.resnet101(False)

iter_cnt = 100
memory_snap = memory_usage(
    (
        forward,
        (net,),
        {
            "use_mkldnn": args.mkldnn,
            "iteration": iter_cnt,
            "weight_cache": args.cache,
            "size": args.size,
        },
    )
)
print(f"Size: {args.size}, memory_usage: {max(memory_snap)}")
</code></pre>
<p>I ran:</p>
<pre data-code-wrap="bash"><code class="lang-bash">python3 test_memory.py 224 10
python3 test_memory.py 256 10
python3 test_memory.py 512 10
</code></pre>
<p>Could someone explain me why this behavior is happening ?</p>
<p>Thanks.</p>