Method for efficiently transferring non-autograd tensors to CPU from GPU?

henryald: <p>My code generates high dimensional Tensor objects on the GPU and subsequently stores these tensors for computation later on. However, there is a large possibility of the GPU running out of memory to store these samples. I have seen the following solution <a href="https://discuss.pytorch.org/t/modifying-forward-backward-pass/169687/2" class="inline-onebox">Modifying forward/backward pass - #2 by ptrblck</a> a few times on the forum, but this solution seems specific to autograd objects, where a graph is involved. However, I am not tracking any gradients.</p>
<p>I could call <code>.cpu()</code> on every sample I store, but this is extremely slow e.g. my tests show the addition of <code>.cpu()</code> slows the code down by up to 10x. I could implement my own cache, which stores a pre-defined number of samples on the GPU before moving them to the CPU, but I was wondering if PyTorch already has its own framework for this.</p>
<p>Thanks.</p>

ptrblck: <aside class="quote no-group" data-username="henryald" data-post="1" data-topic="218032">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/838e76/48.png" class="avatar"> henryald:</div>
<blockquote>
<p>I could call <code>.cpu()</code> on every sample I store, but this is extremely slow e.g. my tests show the addition of <code>.cpu()</code> slows the code down by up to 10x.</p>
</blockquote>
</aside>
<p>Your profiling code is most likely wrong as you might not be synchronizing the code before starting and stopping host timers. The <code>cpu()</code> call will implicitly synchronize your code and will thus accumulate the runtime of already scheduled kernels.</p>

henryald: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="218032">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Your profiling code is most likely wrong as you might not be synchronizing the code before starting and stopping host timers. The <code>cpu()</code> call will implicitly synchronize your code and will thus accumulate the runtime of already scheduled kernels.</p>
</blockquote>
</aside>
<p>Sorry, I’m not sure I follow. The profiling is as follows:</p>
<pre><code class="lang-auto">import torch
import time

gpu = "mps"

def test_gpu(n_samples, *tensor_size):
    obj = torch.zeros((int(n_samples), *tensor_size), device=gpu)
    a = torch.randn(*tensor_size, device=gpu)
    start = time.time()
    for i in range(int(n_samples)):
        obj[i] = a.clone()
    end = time.time()

    print((end - start) / n_samples)

def test_cpu(n_samples, *tensor_size):
    obj = torch.zeros((int(n_samples), *tensor_size), device="cpu")
    a = torch.randn(*tensor_size, device=gpu)
    start = time.time()
    for i in range(int(n_samples)):
        obj[i] = a.clone().cpu()
    end = time.time()

    print((end - start) / n_samples)

test_gpu(1e6, 64, 64)
test_cpu(1e6, 64, 64)
</code></pre>
<p>However, the extent to which calling <code>.cpu()</code> slows down the code is system dependent: on <code>mps</code> the code could be 10-100x slower while on a <code>cuda</code> system it is only 2x slower to call <code>.cpu()</code>. But it is slower nonetheless.</p>

ptrblck: <p>Because this call implicitly synchronized your code. Add <code>torch.cuda.synchronize()</code> before starting and stopping host timers to measure the kernel execution time alone.</p>

neoncube: <p><a class="mention" href="/u/henryald">@henryald</a> Alternatively, you can use <code>torch.profiler</code>, which does the syncing for you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=14" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>For example: <a href="https://discuss.pytorch.org/t/model-uses-gpu-but-backwards-doesnt/218667/3" class="inline-onebox">Model() uses GPU but backwards() doesn't - #3 by neoncube</a></p>